3.2.1 Sarathi-Serve:吞吐量-延迟权衡的突破性解决方案
论文信息
- 标题:Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve
- 作者:Amey Agrawal et al. (Microsoft Research India & Georgia Tech)
- 会议:OSDI 2024
- 代码:https://github.com/microsoft/sarathi-serve
问题定义
现有LLM推理系统面临一个根本性的吞吐量-延迟权衡困境:
Prefill优先调度(如 Orca、 vLLM):
- 优先处理新请求的 prefill阶段
- 优势:快速形成大批量 decode,提升整体吞吐量
- 劣势:导致 "生成停顿 "( generation stalls),严重影响TBT(Time-Between-Tokens)延迟
Decode优先调度(如 FasterTransformer):
- 等待所有运行请求完成 decode后才处理新prefill
- 优势:保证低TBT延迟
- 劣势:batch size小,吞吐量严重受限
核心问题:当batch中混入长prompt的prefill时,整个batch的执行时间会急剧增加,导致decode请求的TBT延迟飙升。实验表明,vLLM中的生成停顿可以持续数秒之久。
核心思想:Chunked Prefill + Stall-free调度
Sarathi-Serve提出了两个核心创新
- Chunked Prefill(分块预填充)
将长prompt的prefill阶段拆分成多个小chunk,每个chunk与decode请求一起batch执行:
传统方式:
[Prefill_A(1024 tokens)] → [Decode_A] → [Decode_A] → ...
↑
长延迟阻塞
Chunked Prefill :
[Chunk_A1(256 tokens) + Decode_B + Decode_C] →
[Chunk_A2(256 tokens) + Decode_B + Decode_C] →
[Chunk_A3(256 tokens) + Decode_B + Decode_C] → ...关键洞察:由于decode的计算强度极低,处理1个decode token的时间几乎等于处理128个prefill token的时间。因此,可以将prefill chunk与decode请求"搭便车"(piggyback)执行,而不会显著增加迭代延迟。
- Stall-free Scheduling(无停顿调度)
Sarathi-Serve的调度器遵循以下原则: 1. 首先打包所有运行中的 decode请求 2. 然后包含部分完成的prefill chunk 3. 最后才接纳新请求,并限制prefill chunk的大小
通过设置token budget(每轮迭代的最大token数),确保每个迭代的执行时间有明确上界,从而避免生成停顿。
技术细节
混合Batch的算术强度分析:
设batch中有 \(B\) 个decode请求和1个大小为 \(P\) 的prefill chunk:
- Decode-only batch的算术强度: \(AI_{decode} = \frac{2 \cdot B \cdot d^2}{B \cdot d} = 2d\)(内存瓶颈)
- Hybrid batch的算术强度:\(AI_{hybrid} = \frac{2 \cdot (B + P) \cdot d^2}{(B + P) \cdot d} =2d\)(当 \(P \gg B\) 时)
当 \(P\) 足够大时,hybrid batch可以达到与纯prefill相近的算术强度,从而充分利用GPU计算资源。
Pipeline Parallelism 优化:
Chunked Prefill还有一个重要副作用:使得每个迭代的计算负载更加均匀,从而显著减少pipeline bubble。在Falcon-180B的8卡部署中,这一优化带来了额外的性能提升。
实验结果
| 模型配置 | 基线系统 | 容量提升 |
|---|---|---|
| Mistral-7B (单A100) | vLLM | 2.6× |
| Yi-34B (2×A100 TP) | vLLM | 3.7× |
| Falcon-180B (8×A100 PP+TP) | vLLM | 5.6-6.9× |
关键发现:
- 在保持相同 TBT 延迟 SLO 的前提下, Sarathi-Serve 显著提升系统容量
- Chunk size是一个可调参数:更大的 chunk提升吞吐量但增加 TTFT
- Pipeline并行场景下收益更大,因为减少了bubble
对Mooncake的启发
Sarathi-Serve的 Chunked Prefill思想被 Mooncake继承和发展。在 Mooncake中, chunkedprefill不仅用于单节点优化,还被扩展到跨节点的PD分离场景中,实现了更细粒度的资源调度。
3.2.2 Parrot:基于语义变量的LLM应用优化
论文信息
- 标题:Parrot: Efficient Serving of LLM-based Applications with Semantic Variable
- 会议:OSDI 2024
问题定义
现有LLM推理服务存在以下问题: 1. 请求级优化局限:公共LLM API以单个请求为单位进行优化,缺乏对应用级语义的理解 2. 请求间依赖未知:无法获知应用内部的请求依赖关系,导致调度效率低下 3. 调度偏好差异:同一应用内的不同请求可能有不同的优化目标(如Map请求关注吞吐量,Reduce请求关注延迟)
核心思想:语义变量抽象
Parrot提出了语义变量(Semantic Variable)这一统一抽象,用于向LLM服务暴露应用的请求依赖信息:
# 传统方式:每次调用都是独立的
response1 = llm.generate(prompt1)
response2 = llm.generate(prompt2)
result = combine(response1, response2)
# Parrot方式:通过语义变量表达依赖
sv1 = parrot.variable("doc_chunk_1")
sv2 = parrot.variable("doc_chunk_2")
summary1 = llm.summarize(sv1) # Map阶段
summary2 = llm.summarize(sv2) # Map阶段
final = llm.combine(summary1, summary2) # Reduce阶段通过语义变量, Parrot能够: 1. 分析请求依赖图:识别可并行化的请求 2. 优化调度策略:根据请求类型选择不同的批处理和调度策略 3. 跨请求优化:复用相同语义变量的KVCache
应用场景
文档摘要应用:
- Map阶段:并行处理文档的各个chunk,关注吞吐量
- Reduce阶段:合并摘要结果,关注延迟
多轮对话应用:
- 识别对话历史中的语义变量
- 跨会话复用系统prompt的KV Cache
3.2.3 InfiniGen:动态KV Cache管理框架
论文信息
- 标题:InfiniGen: Efficient Generative Inference of Large Language Models withDynamic KV Cache Management
- 作者:Wonbeom Lee et al. (Seoul National University)
- 会议:OSDI 2024
- 论文:https://arxiv.org/abs/2406.19707
问题定义
长文本生成场景面临严峻的KV Cache内存瓶颈
对于长度为 \(L\) 的序列, KV Cache 的内存占用为:
\(Memory_{KV} = 2 \cdot L \cdot d_{model} \cdot n_{layers} \cdot \text{batch\_size} \cdot \text{bytes\_per\_param}\)
例如,对于LLaMA-65B模型,batch size=1,序列长度=100K时:
- 每层KV Cache:\(2\times 100K \times 8192 \times 2B = 3.2GB\)
- 80 层总 KV Cache: \(256GB\) 远超单卡显存
现有offloading方案的问题:
- 需要频繁在GPU和CPU内存之间传输完整的KV Cache
- 传输开销成为新的瓶颈
核心思想:重要性感知的KV Cache预取
InfiniGen的核心洞察:并非所有token的KV Cache都同等重要
通过分析attention机制,InfiniGen发现: 1. 只有少数token对后续预测有重要影响 2. 重要token的集合在不同attention head之间存在高度相似性
技术方案:
- 重要Token识别
通过在当前层进行 " 最小化预演 " ( minimal rehearsal ):
- 使用当前层的输入和部分query/key权重
- 预测下一层 attention中哪些 token会获得高 attention score
- 只预取这些重要token的KV Cache
- I/O优化的KV管理
传统Offloading:
每轮都需要加载全部KV Cache → 高I/O开销
InfiniGen:
预识别重要token → 只加载重要KV → I/O开销降低实验结果
| 模型 | 序列长度 | 基线 | InfiniGen | 加速比 |
|---|---|---|---|---|
| LLaMA-2-7B | 128K | DeepSpeed | InfiniGen | 3.0× |
| LLaMA-2-13B | 128K | DeepSpeed | InfiniGen | 2.5× |
与Mooncake的关联
InfiniGen 的重要性感知思想与 Mooncake 的 KV Cache 分层存储策略有相似之处。Mooncake通过全局调度器识别热数据,将高频访问的 KV Cache保留在快速存储层,低频数据下沉到慢速层,实现类似的I/O优化效果。