2308.16369-sarathi-chunked-prefill-decode-maximal-batching
SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills
Sarathi 的核心贡献是把 LLM serving 的低效从“decode 天然慢”重写为一个可调度的数据流问题:prefill 很快进入 compute-saturating 状态,decode 因为逐 token 生成和 KV cache 访问长期 memory-bound;如果把一个长 prefill 切成多个 compute-sized chunk,再让 decodes 搭在每个 prefill chunk 的 linear operations 上,就可以复用已加载的模型权重、降低 decode 的边际成本,并让 pipeline-parallel inference 的 microbatch 变得更均匀。
Source
- Title: SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills
- arXiv: https://arxiv.org/abs/2308.16369
- HTML: https://ar5iv.labs.arxiv.org/html/2308.16369
- PDF: https://arxiv.org/pdf/2308.16369
- TeX Source: https://arxiv.org/e-print/2308.16369
- DOI: https://doi.org/10.48550/arXiv.2308.16369
- Follow-up: Sarathi-Serve OSDI 2024 https://www.usenix.org/conference/osdi24/presentation/agrawal
- Code/Project: 原 Sarathi 论文未给独立代码;后续 Sarathi-Serve 开源于 https://github.com/microsoft/sarathi-serve
- Authors: Amey Agrawal, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S. Gulavani, Ramachandran Ramjee
- Submitted: 2023-08-31
- Current version read: arXiv v1, submitted Thu, 31 Aug 2023 00:03:02 UTC
- Subjects: Machine Learning (
cs.LG); Distributed, Parallel, and Cluster Computing (cs.DC)
作者与关系
- Amey Agrawal: Georgia Institute of Technology.
- Ashish Panwar: Microsoft Research India.
- Jayashree Mohan: Microsoft Research India.
- Nipun Kwatra: Microsoft Research India.
- Bhargav S. Gulavani: Microsoft Research India.
- Ramachandran Ramjee: Microsoft Research India.
阅读目标与判断边界
本笔记关注:
- 为什么 prefill 与 decode 在 GPU 上呈现完全不同的 arithmetic intensity 和 batch scaling。
- chunked-prefills 与 decode-maximal batching 如何把 decode 从低利用率路径转成 piggybacked marginal work。
- P:D ratio、chunk size、batch size 和 tile quantization 如何共同决定收益。
- Sarathi 与 Orca/vLLM/TGI、FlashAttention、Parrot、Span Query、RL rollout consistency 的关系。
判断边界:
- 论文目标是 inference throughput 和 pipeline utilization,未系统优化 tail latency、queueing delay、公平性、SLO 或多租户隔离。
- 论文实验以 LLaMA-13B/A6000、LLaMA-33B/A100 和 GPT-3 64 A100 simulation 为主;当前 vLLM/SGLang/TensorRT-LLM 的 chunked prefill、continuous batching、PD disaggregation 和 attention backend 已进一步演化。
- Sarathi 优化 decode 的 linear operation 权重复用,对 attention 的 memory-bound 部分帮助有限;上下文越长,attention/KV cache 成本越会限制总收益。
- follow-up Sarathi-Serve 处理 throughput-latency tradeoff,是独立 OSDI 2024 系统,应作为后续阅读节点。
论文脉络
1. 研究问题、背景和价值
LLM inference 分为两个阶段:
- prefill:一次处理完整 prompt,输入 tensor 形状近似为
,大多数操作是 matrix-matrix multiplication。 - decode:每轮只处理新生成的一个 token,输入 tensor 形状近似为
,attention 还需要读取该请求历史 token 的 KV cache。
这两个阶段在 GPU 上的行为差异很大。prefill 只要 token 数足够,就能把 GEMM 做大,快速达到高 GPU 利用率;decode 每轮 token 数少,模型权重和 KV cache 读写占主导,实际更接近 memory-bound vector-matrix workload。
论文用 LLaMA-13B on A6000 profiling 给出一个直接信号:sequence length 为 1024 时,decode per-token cost 在 batch size 1、2、18 下分别约为 prefill 的
这个问题重要,因为 serving 成本大量花在 decode 上。请求只做一次 prefill,却要做很多轮 decode;decode 每轮都要访问模型权重和 KV cache。若只靠扩大 batch 或 model parallelism,系统会受到显存、通信、tail latency 和 pipeline bubble 限制。Sarathi 的价值在于它把 prefill 的高算力利用率转化成 decode 的载体,让 decode 作为 prefill 的边际增量执行。
2. 已有解决方案与不足
已有 serving 调度大致有三类。
第一类是 request-level batching,例如 FasterTransformer 风格的一批请求一起跑到完成。它实现简单,但短请求会等待长请求,prefill-only 和 decode-only 阶段分离,decode 仍低效。
第二类是 iteration-level scheduling,例如 Orca、vLLM、HuggingFace TGI。请求可以在每个 iteration 进入或退出 batch,减少 padding 和等待。但这些系统通常把一个新请求的整个 prompt 作为一次完整 prefill 插入 batch。这样会带来两个问题:full prefill 太大,容易拖慢正在 decode 的请求;full prefill 数量有限,能 piggyback 的 decode 覆盖率不稳定。
第三类是 model parallelism。Tensor parallelism 可以让更大 batch 放入显存,但层内 all-reduce 通信在 critical path 上,通常更适合同节点 NVLink;pipeline parallelism 更适合跨节点大模型部署,通信只在 stage 间传 activation,但 microbatch 的执行时间不均会造成 pipeline bubbles。LLM inference 中 prefill 长度、decode 长度、KV cache length 都会变化,bubble 依然明显。
因此,论文留下的空白是:如何构造一串计算量更均匀、能吃满 GPU、还能覆盖大量 decode iteration 的工作单元。
3. 作者可能的思考路径
作者大概率先从 profile 看到两个现象:prefill 很快 saturate,decode 在可行 batch size 下长期 memory-bound;同时 pipeline parallelism 下,bubble 来自 microbatch 执行时间差异。若只从“decode 很慢”出发,直觉方案是增大 decode batch,但 KV cache 容量限制会挡住这条路。若只从“prefill 很快”出发,直觉方案是让更多 prefill 进入 batch,但真实请求里每个 request 只有一次 prefill,decode 轮数更多,prefill 数量不足。
下一步自然是把 prefill 这个稀缺的大块计算拆成多个小块。只要 chunk 仍然足够大,可以保持 matrix-matrix multiplication 的效率;又因为一个 prompt 可以拆成多个 chunk,一个请求就能提供多次“可搭载 decode 的高利用率事件”。这样,decode 不再需要单独把权重从 HBM 加载一遍,而是和 prefill chunk 合并进同一次 linear operation。
最后,pipeline bubble 也可以用同一个思路解释:让每个 microbatch 都包含一个大小接近的 prefill chunk 和尽可能多的 decodes,microbatch compute time 更均匀,stage 间等待下降。这个思路把单 GPU decode efficiency 和多 GPU PP bubble reduction 合并成一个调度 primitive。
4. 核心假设或切入点
Sarathi 的核心假设是:decode 的低效主要来自 linear operations 中模型权重重复加载和小 token batch 的低 arithmetic intensity;attention/KV cache 部分仍然要单独处理,但 linear operations 在 forward pass 中足够重要,因此 piggyback decodes 可以显著降低 decode 的边际成本。
第二个假设是:prefill 可以被切成多个 chunk,并通过正确 attention mask 保持与完整 prefill 数学等价。对第
5. 方法 / 系统 / 理论框架
Chunked-prefills
假设 prompt 长度为
- 当前 chunk 生成自己的
。 - attention mask 允许当前 chunk 的 query 访问所有已处理 prompt token 的
和当前 chunk 内满足 causal 顺序的 token。 - 当前 chunk 的
写入 KV cache,供后续 prefill chunk 和后续 decode 使用。
chunking 的代价来自两处:
- chunk 太小会降低 arithmetic intensity,使 prefill GEMM 不够大。
- 后续 chunk 需要重复读取此前 chunk 的 KV cache;第一个 chunk 的 KV 会被读
次,第二个 chunk 读 次,依此类推。
论文的经验判断是,attention 在整体 forward pass 中占比相对有限,合理 chunk size 可以控制额外开销。LLaMA-13B on A6000 中,chunk size 256/512 能把端到端 prefill computation loss 限制在约 20%/10% 内。
Decode-maximal batching
Sarathi 每个 batch 放入一个 prefill chunk,再用 decodes 填满剩余 slot。可用 decode batch size 受 GPU 显存限制:
其中
关键实现是:linear operations 把 prefill token 和 decode token 合并成一个更大的 matrix multiplication;attention 仍将 prefill attention 与 decode attention 分开执行。这样做的原因是 linear layer 使用同一组模型权重。prefill chunk 已经要把权重从 HBM 取入计算路径,decode token 可以在同一次权重加载下完成计算,decode 的边际成本下降。
论文给的 LLaMA-13B/A6000 例子很直观:
| Batch type | Linear time | Attention time | Total time | Decode cost |
|---|---|---|---|---|
| prefill-only, 1024 tokens x4 | 224.8 ms | 10.0 ms | 234.8 ms | - |
| decode-only, batch 4, seq 1024 | 44.28 ms | 5.68 ms | 49.96 ms | 12.49 ms/token |
| 1021 prefill + 3 decodes | 223.2 ms | 15.2 ms | 238.4 ms | 1.2 ms/token |
这说明 piggyback decode 后,decode per-token marginal cost 可以降低约一个数量级。
P:D ratio 与理想 chunk size
论文用 P:D ratio 表示 prefill token 数与 decode token 数之比。若 batch size 为
因此峰值通常出现在:
这个公式解释了 Sarathi 的主要调参逻辑:
- chunk 更小:prefill chunk 数更多,能覆盖更多 decode iterations,但 prefill 本身效率下降。
- chunk 更大:prefill 更高效,但能 piggyback 的 decode 更少。
- batch 更大:每个 chunk 可搭载更多 decodes,理想 P:D 变小。
此外,论文强调 tile quantization。GPU matmul 会按 tile 切分;当 relevant matrix dimension 不是 tile size 的倍数,thread block 会做额外无效计算。论文实验中 128 到 256 tokens 只增加 27% runtime,但 256 到 257 tokens 增加 32% runtime。因此实际 chunk size 要让:
尽量对齐 tile size。论文给出经验例子:若 nominal chunk size 为 256、tile size 为 128、maximum batch size 为
6. 结论链条
Sarathi 的论证链是:
- profile 证明 prefill 和 decode 的 GPU 行为差异很大,decode 在实际 batch range 内 memory-bound。
- chunked-prefills 把单次 prefill 扩展成多次 compute-saturating events,并通过 causal attention mask 保持数学等价。
- decode-maximal batching 把 decode token 放进 prefill chunk 的 linear operations,复用模型权重加载,降低 decode 边际成本。
- P:D ratio 和 chunk size 决定 prefill chunk 数与 decode iteration 数是否匹配。
- 均匀 compute units 同时降低 pipeline parallelism 的 bubble。
- 单 GPU和多 GPU实验都显示收益,但端到端 throughput gain 明显低于 decode-only speedup,因为 prefill 和 attention 部分仍有成本。
关键实验/定理
结果 1: Decode speedup
- 设置:LLaMA-13B on A6000,chunk size 256,vary batch size 与 sequence length。
- 指标:decode throughput speedup,以 piggyback batch 相对 prefill-only batch 的边际 runtime 估计 decode cost。
- 结果:decode throughput improvement 为
到 。 - 解读:收益来自 linear operations 的权重复用;batch size 变大或 sequence length 变长时,baseline decode 更有效且 attention 成本增加,speedup 下降。
结果 2: 单 GPU端到端 throughput
- 设置:LLaMA-13B on A6000;LLaMA-33B on A100;sequence length 1K/2K/3K;chunk size 256。
- 指标:decode speedup 与 end-to-end throughput gain。
- 结果:
| Model / GPU | Sequence | Batch | P:D | Decode speedup | Throughput gain |
|---|---|---|---|---|---|
| LLaMA-13B / A6000 | 1K | 6 | 50:1 | ||
| LLaMA-13B / A6000 | 2K | 6 | 50:1 | ||
| LLaMA-33B / A100 | 1K | 10 | 28:1 | ||
| LLaMA-33B / A100 | 2K | 5 | 63:1 |
- 解读:decode speedup 可以很大,端到端收益在 1.14-1.33x 区间,原因是 Sarathi 主要优化 decodes,prefill 和 attention 仍保留成本。
结果 3: P:D ratio 峰值公式
- 设置:vary P:D ratio、sequence length、chunk size、batch size。
- 指标:end-to-end throughput。
- 结果:峰值通常出现在
;例如 、 时,峰值 P:D 约为 ,论文观测到 1K setting 下 P:D=14 左右达到 peak。 - 解读:收益要求 prefill chunks 与 decode iterations 数量匹配。P:D 太低时 prefill chunk 不够,decode 剩余部分要单独跑;P:D 太高时 decode 不够,prefill chunk 缺少可 piggyback 工作。
结果 4: Orca / iteration-level scheduling 对比
- 设置:Orca best-case 与 worst-case;Sarathi chunk size 256/512;sequence length 1K/2K/3K。
- 指标:overall throughput gain。
- 结果:Orca best-case 在 1K 上约
,Sarathi 在 1K/2K/3K 分别约 、 、 。 - 解读:Orca 的 prefill/decode overlap 来自新请求进入 batch 时的 incidental overlap;Sarathi 把 full prefill 切成多个 chunk,提供更稳定、更细粒度的 piggyback 机会。
结果 5: Pipeline parallelism simulation
- 设置:GPT-3,64 A100,8 servers;8-way TP within node + 8-way PP across nodes;对比 TP+PP Orca-style、TP+PP Sarathi、8 replicas TP-only。sequence length 1K-4K,Zipf
,P:D=10,chunk size 256。模拟器基于 8-GPU A100 DGX empirical profiling,runtime 估计误差在 5% 内。 - 指标:pipeline bubble time per request;10K requests completion time。
- 结果:Sarathi median bubble time per request 降低
;TP-only 比 baseline TP+PP Orca 快 ;Sarathi TP+PP 比 baseline TP+PP 快 ,比 TP-only 快 。 - 解读:Sarathi 的均匀 compute units 让 PP 的大 batch / 大 KV cache 容量优势得以释放,减少跨 stage 等待。
结果 6: Chunking overhead ablation
- 设置:LLaMA-13B on A6000,chunk size 64-512。
- 指标:attention overhead、prefill-only runtime overhead、end-to-end throughput。
- 结果:chunk size 64 带来约
attention overhead 和约 overall prefill time;chunk size 256/512 将端到端 prefill computation loss 控制在约 20%/10%;chunk size 128 即使 prefill 慢超过 ,仍可获得最高约 throughput,因为可 piggyback 更多 decodes。 - 解读:prefill efficiency 与 decode coverage 之间存在真实 tradeoff;chunk size 需要按 workload 和硬件 profile 调。
证据链强度评估
强证据
- prefill/decode arithmetic intensity 和 per-token cost profiling 支撑了问题定义。
- chunked-prefill 的 causal attention mask 保证语义等价,这一点机制清楚。
- 单 GPU实验覆盖 LLaMA-13B/A6000 与 LLaMA-33B/A100,显示方案跨模型/硬件有效。
- pipeline simulation 的输入来自 empirical profiling,且作者报告与 8-GPU A100 DGX empirical value 误差在 5% 内。
中等强度证据
- P:D ratio 公式解释力强,但真实 workload 中 P、D、arrival pattern、SLO、tenant mix 会动态变化,固定 chunk size 未必长期最优。
- Orca 对比采用 best-case/worst-case,能说明 full prefill overlap 的限制,但现代 vLLM/SGLang chunked prefill 和 scheduler 已经继续演化,需要新 baseline。
- Sarathi-Serve follow-up 说明该思想继续发展为低延迟高吞吐 scheduler,但原 Sarathi 与 Sarathi-Serve 的贡献边界需要分开。
需要谨慎的推论
- 论文的端到端收益不能直接外推到 10K-100K 或 1M context;此时 attention/KV cache 读写会更重。
- Decode speedup 高达
,但总体成本下降通常在 25% 量级;部署价值取决于 workload 中 decode 占比和成本结构。 - 若用于 RL rollout 或评测,chunking/mixed batching 可能改变 kernel path、precision path、batch composition 和 nondeterminism,需要与 2025-09-10 / 2605.14220 一起审计。
OpenReview / 审稿意见吸收
- Venue status: 当前档案未记录公开 peer-review 状态。
- Public reviews: 当前档案未记录可可靠匹配的 OpenReview / ARR / 会议 reviewer comments。
- Ratings / confidence: 无公开评分可用于校准。
- Reviewer consensus: 暂无。
- Main criticisms: 暂无公开 reviewer 质疑可引用;可信度主要由论文、技术报告、项目证据和本地一致性检查决定。
- Author response: 暂无公开 rebuttal 记录。
- 对本文可信度的影响: 按未完成公开审稿吸收处理,结论需要依赖实验设置、baseline 强度、复现证据和跨论文一致性校准。
本地讨论补充
1. 讨论收敛点
- Sarathi 和 FlashAttention 处在不同层级:FlashAttention 优化单个 attention kernel 的 HBM/SRAM IO;Sarathi 优化 serving scheduler 形成的 batch composition 与 linear operation 权重复用。
- Sarathi 和 Parrot/Span Query 也处在不同层级:Parrot/Span Query 要求上层应用暴露结构;Sarathi 即使只看到普通 prompt/decode 阶段,也能通过 chunked prefill 改变执行粒度。
- Sarathi 适合放在 LLM serving 图谱里的 “prefill/decode scheduling” 节点,作为后续 Sarathi-Serve、POD-Attention、PD disaggregation、chunked prefill、continuous batching 的早期基础材料。
2. 修正后的理解
- “decode 很慢”更精确地说是:decode 的可并行 token 数小,linear operations 反复读取模型权重,attention 反复读取 KV cache;在可行 batch size 下,GPU compute units 很难被吃满。
- chunked-prefill 的重点在于把一次 prefill 转换成多次可与 decodes 对齐的 compute-saturating work item;它本身并不减少 prefill FLOPs。
- decode-maximal batching 的收益来自 linear operation 的融合;attention 仍分开执行,因此长上下文下 attention/KV cache 会逐步成为上限。
3. 后续复验指标
- 在现代 vLLM/SGLang/TensorRT-LLM 上比较 Sarathi-style chunked prefill、continuous batching、PD disaggregation、prefix cache、speculative decoding 的组合收益。
- 对真实 request trace 统计 P:D ratio 分布、prompt length distribution、output length distribution、arrival burst、SLO violation 与 fairness。
- 分析 chunk size 对 TTFT、ITL、TPOT、throughput、GPU memory、KV cache hit、tail latency 的联合影响。
- 在 RL rollout workload 中检查 mixed batching 是否造成 logprob mismatch、batch-dependent nondeterminism 或 trainer/rollout backend 差异。
主要启发
- LLM serving 的优化对象不只是单次 forward kernel,也包括“哪些 token 在同一个 batch 里共同使用权重”。
- prefill 和 decode 的差异可以转化成互补:prefill 提供高 arithmetic intensity,decode 提供长期低效的尾部工作;合理混合后,两者可以共享一次权重读取。
- chunk size 是 serving scheduler 的核心控制变量,连接了硬件 tile、prefill efficiency、decode coverage、batch size、P:D ratio 和 latency。
- Pipeline parallelism 在 inference 中的 bubble 不只来自传统 warmup/drain,也来自 LLM 请求内部 prefill/decode/length 不均;均匀 compute unit 能让 PP 的显存优势重新有用。
- 对后续系统论文阅读,应该区分四层:kernel 层 attention/linear 实现,KV cache memory 管理层,batch/scheduler 层,应用 DAG / semantic structure 层。
局限
- 论文主要优化 throughput,真实 serving 还需要同时处理 TTFT、ITL、queueing delay、公平性、SLO 和多租户隔离。
- chunk size 需要 workload、模型和硬件 profile;P:D ratio 未知或动态变化时,固定 chunk size 可能偏离最优。
- 除 simulation 外,论文简化假设同一 batch 中请求具有相同 prefill/decode token 数;真实流量长度差异更复杂。
- 实验上下文长度最高到 3K,P:D ratio 范围为 1-200;长上下文进入 10K-100K 后,attention quadratic cost 和 KV cache memory traffic 会改变瓶颈结构。
- 实现基于 nanoGPT 和当时 xFormers attention;现代 serving engine 的 baseline 已加入 paged KV cache、chunked prefill、continuous batching、FlashAttention-2/3、FlashInfer、PD disaggregation 等。
跨论文关系
- 与 Grape:Sarathi 把一个请求的 prefill 切成 chunk,与 decode 组成 decode-maximal batch;Grape 将相同执行原语扩展到相依 workflow task,让上游 decode chunk 直接成为下游增量 prefill 的输入,并在 SLO 约束下调度这些微任务。两者分别提供 chunk primitive 与跨 task dataflow lowering。
- 与 2205.14135 / 2307.08691:FlashAttention 系列减少 attention kernel 的 HBM traffic 和提高 SM/warp utilization;Sarathi 把 prefill/decode 组织成更适合 GPU 的 batch,并让 decode 复用 prefill linear operations 的权重读取。两者共同说明 LLM systems 需要同时看 memory hierarchy 与 workload composition。
- 与 2405.19888:Parrot 从应用 DAG、Semantic Variable 和 shared prefix 暴露优化机会;Sarathi 从普通 inference request 的 prefill/decode 阶段切分优化机会。二者都把 serving 从“单请求黑盒 completion”推进到结构化执行。
- 与 2511.02749:Span Query 把 cache locality 和 attention locality 交给 expression-tree optimizer;Sarathi 处理 prefill/decode locality 与模型权重加载 locality。二者都属于 runtime 通过更细粒度输入结构换取系统效率。
- 与 HybridFlow、verl 官方仓库 和 slime 官方仓库:三者编排 actor、rollout、reward 和 policy update 的训练-生成 dataflow;Sarathi 处理 rollout / inference engine 内部的 prefill/decode scheduling。两层组合时需要同时记录吞吐、长尾、batch composition 与 logprob consistency。
- 与 2025-09-10 和 2605.14220:Sarathi-style mixed batching 在 production inference 中是性能优化;若进入 RL rollout,则要额外记录 batch composition、kernel path、precision 和 logprob 计算路径,避免训练侧把 serving optimization 当成稳定行为分布。
- 跨论文关系定位:记录 Sarathi、Chunked Prefill 与 Prefill/Decode Scheduling,并连接 LLM Serving、FlashAttention、Parrot/Span Query、Inference Determinism 与 TIM。
Reference Intake Brief
Target
- Intended target system: 新增论文笔记、更新
content/utility/papers-index.md、更新data/authors.json。 - Existing related assets: 2205.14135、2307.08691、2405.19888、2511.02749、2025-09-10、2605.14220。
- Proposed form: 新建独立 Markdown 文档,并在跨论文关系中补充 serving scheduler 主题。
Reusable Elements
- Prefill/decode 两阶段的 arithmetic intensity 与 memory-bound/compute-bound 区分。
- Chunked-prefill + decode-maximal batching 的 batch construction 公式和 P:D 公式。
- Sarathi 与 Orca/vLLM/TGI、FlashAttention、Parrot/Span Query、RL rollout consistency 的跨层定位。
Risks
- Copyright/over-copying: 使用论文事实、公式和数值结果,正文为本地重述;避免长段复制。
- Unsourced or unverifiable claims: 作者与机构来自 arXiv、ar5iv、MSR 页面、个人主页、GitHub、USENIX;Bhargav 的公开主页证据较弱,已标注。
- Tone/brand mismatch: 保持系统论文分析语气,避免把 follow-up Sarathi-Serve 的贡献混同为原论文贡献。
- Safety/compliance issues: 无直接双用途安全操作细节。
- Overlap with existing assets: 与 FlashAttention、Parrot、Span Query 存在主题关系,但层级不同,值得独立归档。
Skipped
| Material | Reason |
|---|---|
| Sarathi-Serve full paper | 后续系统,和原 Sarathi 关系强但贡献不同,应单独归档 |
| POD-Attention / QoServe / SUTRADHARA | 同属 MSR India serving line,当前任务只处理 2308.16369 |
| 完整复现实验 | 需要现代 serving engine、GPU 和 workload trace,超出本轮阅读 |
Recommendation
Decision: merge
Why: Sarathi 是当前 LLM serving 图谱中 prefill/decode scheduling 的基础节点,能够连接 kernel-level attention efficiency、serving-level continuous batching、pipeline parallelism、后续 Sarathi-Serve 和 RL rollout consistency。