Paper Reading: LoongServe

Paper Link: LoongServe: Efficiently Serving Long-context Large Language Models with Elastic Sequence Parallelism

Background and motivation

1.1 背景

LLM Serving的普遍存在的问题:

  • LLM请求的模式是高度动态变化的,部分的Context可能会很长
  • Prefill和Decode对于资源的需求也是不同的,Prefill是Computation bound,Decode是Memory bound

关于如何serve 长文本:

  • Sequence Parallelism:将每一个Request切分sequence
  • Tensor Parallelism:切分WQ、WK、WV矩阵(是一种Model Parallelism)

Figure 1: 现有的Sequence Parallelism

img

补充知识:Tensor Parallelism and Sequence Parallelism: Detailed Analysis

Tensor Parallelism:

img

Sequence Parallelism:

img

1.2 Existing Solution

为了应对动态变化的sequence length和不同的资源需求,通常采用GPU分组的方式:

  • 每组GPU部署一个LLM实例
  • 每组GPU采用不同的并行策略(sequence parallelism,tensor parallelism)或者讲prefill & decode分开

1.3 存在的问题

  1. SP: 只支持Prefill阶段;应对固定DoP的training场景,静态分组策略不能满足动态的资源需求
  2. 在进行scale时,迁移KV cache带来额外的通信开销
  3. 不同group的GPU里的内存不能被gather起来,导致memory fragmentation

Figure 2: Tensor Parallelism下,Batch Size和sequence length对每次iteration latency的影响

  • 处理 100K input tokens 比处理 1K input tokens 慢 105.97 倍
  • 由于Prefill Phase是compute-intensive,因此增加 DoP 可以显著加速;相反,对于计算量相对较轻的Decode Phase,由于额外的通信开销,较大的 DoP 导致性能提升较小
  • 因此,静态并行策略不能在所有情况下都有效

img

Figure 3: 将 SP 与现有并行策略(如TP)结合使用不会带来额外的开销,甚至对不同context sequence和decoding phase请求都能获得更好的性能

img

系统设计

核心思想:

  • 增加对于Decode Phase的Sequence Parallellism支持
  • 在不同的phase之间灵活管理KV Cache
  • 动态分发请求(由Global Manager完成)

因此系统架构分为两部分:

  1. Elastic Instances
  2. Global Manager

img

2.1 Elastic Instances

Loong Serve提出了zero-overhead ESP (Elastic SP) 机制

img

#1 Elastic Scale Down

在Prefill Phase结束之后,由于Decode Phase对于计算资源的需求少,所以需要进行Scale Down

Existing Solution:Reactive migration

  • 在Prefill阶段结束,将 key-value tensors迁移到新的Parallel Group (𝑅′) 中,通信开销严重
  • Memory fragmentation问题:𝑅′ 里的每一个实例都需要至少 𝑂 (𝑏𝑙𝑠ℎ/𝑑) 的GPU内存才放得下新生成的 KV tensors

LoongServe Solution:Proactive migration reuses existing communication results.

  • 三个Instance缩到两个,前四个Token放在Instance 1,其他的Token放在Instance 2
  • 在Prefill阶段进行P2P通信从neighbor获取KV tensor时候,不止是计算attention然后发给下一个Instance,还要根据预先生成的Scale Down Plan选择性地保存KV tensor (也就是KV Cache)

img

#2 Elastic Scale Up

在Decode阶段,KV Tensor可能会变长,超出当前 Instance 的 KV Cache Pool 容量,需要进行Scale up

Existing solution (Tensor Parallelism,FlashDecoding)

  • 当Memory不够的时候,需要把Batch里的一些request迁移到其他instance,带来KV Cache migration开销

LoongServe Solution:

首先提出了single-master distributed decoding:

  • Extend Sequence Parallelism to Decode Phase,允许多个Instance参与单个Batch的Decode计算
  • 在计算Attention的时候,每一个Instance计算本次的Local Attention,然后发给Master做Reduce
  • FFN layer在单卡(master)上做,生成下一个Token
  • 增加新的机器时,只需要让这个新的机器作为Master,不用迁移KV Cache

进一步优化:multi-master distributed decoding

  • 当Batch size过大时,master的压力比较大
  • 一个Batch内的不同Request,可以在不同的GPU上计算FFN layer(即生成新的Token),即不同的Reqeust有不同的Master

img

2.2 Global Manager

将scheduling problem 解构成四个子问题:

#1 dispatching

  • Global Manager按照先进先出(FCFS)的顺序扫描pending queue中的请求
  • 考虑GPU内存限制:确保为新请求保留足够的未使用slot
  • 考虑GPU计算限制:使用分析模型和SIB中的分析结果来评估对现有Decode阶段请求的影响,并决定是否将新请求添加到当前iteration中

#2 elastic instance allocation

  • 主要考虑减少干扰:尽量避免Prefill阶段和Decode阶段之间的干扰,同时最大化GPU计算的效率
  • Global Manager首先为选定的请求分配空闲的实例
  • 如果空闲实例的未使用KV Cache Slot不足,可能会抢占一些实例以获得足够的slot

#3 batching

  • Global Manager根据请求的序列长度对请求进行排序,并为它们分配不同的DoP(Degree of Parallelism)
  • Key insight是具有相似sequence length的请求应该被batch在一起
  • 使用动态规划来优化Batching策略,目的是最小化prefill阶段的input latency

#4 elastic scaling plan generation

  • 对于proactive scaling down,Global Manager设置模型并行度为解码阶段的最小最佳DoP,并在运行时根据需要调整DoP
  • 对于scaling up,当GPU计算或内存不足时,Global Manager会增加新的实例到parallelism group中,使用multi-master decoding来减少内存碎片或执行时间

Evaluation

4.1 Setup

Model: LWM-1M-Text model

Testbed: 8* NVIDIA A800 80GB GPUs, 128 CPUs, 2048 GB of host memory, and four 200 Gbps InfiniBand NICs

Workloads: Generated by a Poisson process, samples from real-world datasets

Baselines: vLLM, DeepSpeed-MII, LightLLM w/ SplitFuse, DistServe

Metrics: Focus on the serving throughput.

  • Normalized per-token latency: the mean of requests’ end-to-end latency divided by their sequence lengths
  • Normalized input latency: the mean of prefill phase time divided by the input lengths
  • Normalized output latency: the mean of decoding phase time divided by the output lengths

4.2 End-to-end performance

  • 对于Prefill Phase:通过设置不同的DoP避免阻塞住Pending的Prefill Phase,取得了性能提升
  • 对于Decode Phase:使用不同的Instance服务不同的Phase,因此Decode Phase不会被Prefill Phase干扰 => Low output latency

img

4.3 Multi-node performance

16 GPU的集群,将ESP设为8

Total throughput and input throughput:

  • 1.86× and 1.72× compared to vLLM
  • 3.37× and 3.11× compared to LightLLM w/ SplitFuse,

img

消融研究

5.1 ESP的有效性分析

四种Baseline:

  • full LoongServe
  • traditional tensor parallelism, i.e., LoongServe w/o ESP (TP=8)
  • static hybrid parallelism, i.e., LoongServe w/o ESP (TP=2, SP=4)
  • parallelism with replication, i.e., LoongServe w/o ESP (TP=2) × 4

证明了只引入某一项技术是不够的,必须结合起来

img

5.2 Scaling Overhead

overhead of scaling down and scaling up under different batch sizes and input lengths by timing the time to forward a batch with and without scaling

  • 对于scaling down:只引入了很小的开销(大约2%)
  • 对于scaling up:对于大的Batch size来说,compute-memory比例很大,将computation分布在多Instance可以显著提升性能

img