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

Background

LLM推理过程[^1]

  1. 模型参数加载到GPU

  2. 在CPU上对prompt分词（tokenizing），并将token tensor传输到GPU

    Tokenization step

  3. 输入token到Transformer，生成第一个token

  4. 将生成的token附加到输入token序列中，将其作为生成第二个token的新输入。然后，重复此过程，直到生成了停止符（EOS，end-of-sequence）或达到最大序列长度

Initiation and decoding phases of the token generation process

  1. 将完成的 tokens 从 GPU 获取到 CPU ，并对它们进行 detokenize（”detokenize“指的是将模型生成的 tokens 序列转换回原始文本或句子的过程。可能包括去除 tokens 之间的空格、添加标点符号、还原缩写等操作，以还原生成文本的自然语言形式。），以获取生成的文本。

Detokenization step

Transformer LLM 的推理过程分为两个阶段，将1-3称作预填充阶段（pre-fill phase）

将4称生成阶段（generation phase）、解码阶段（decoding phase）、自回归阶段（auto-regressive phase）、增量阶段（incremental phase）。

KV Cache[^2]

为方便起见，下面将token的上下文表征（contextual representation）称作token

由于Transformer在Decoder使用了Mask技术，在生成新token时，只能使用之前已生成的token信息。而之前的token在每次迭代中都是相同的，所以当前新生成的token在随后的所有迭代也是相同的，因此存在
冗余计算。

举个🌰，假设在解码阶段新一次迭代，使用了 What color is the sky? The sky作为输入序列生成了is，在上一次迭代中， sky是输入序列最后一个token，因此与该token上下文表征是通过使用序列中所有token的表征生成的，

即What,color,is,the,sky,?,The.sky的key、value向量,

下一次迭代输入序列将是What color is the sky? The sky is，

但由于做了Mask，从sky角度讲，输入序列仍然是What color is the sky? The sky，

因此，sky生成的输出向量与上一次迭代完全相同。

同时，唯一尚未计算的是输入序列中最后一个token is的上下文表征，为此需要得到:

  1. is的query向量
  2. What color is the sky? The sky is的key向量
  3. What color is the sky? The sky is的value向量

对于key和value向量，除了is外，都已经在之前迭代中计算过了，因此可以将其缓存下来，也即KV Cache。

那么，为is计算输出向量将会变得非常简单：

  1. 计算is的query、key、value向量
  2. 从缓存中获取What color is the sky? The sky的key和value向量，并与is的key和value向量连接
  3. 使用is的query向量和所有key向量计算attention score
  4. 使用attention score和所有value向量计算is的输出向量

当我们使用 KV Cache 时，模型输入就不是整个序列，而是最后生成的 token 和 KV Cache

Generation step with KV caching enabled

Pre-fill Phase

将输入的prompts得到q、k、v，存入KV Cache（为Decode阶段缓存）。

这一阶段需要将当前输入序列的KV一次性全部计算出来，因此需要大量计算资源。

Decode Phase

由新生成的token得到q、k、v，计算它与之前所有token的attention。

这一阶段需要从KV Cache中读取前面所有token的Key和Value，因此需要占用大量内存资源。

2 Phase in LLM Inference

Challenges

  1. LLM推理所需要的计算资源及GPU显存占用随着不同请求（输入序列）上下文长度的不同而有很大差异
  2. 即使对于单一请求而言，其在Prefill和Decode两个阶段所需的资源也有很大差异。

Main Idea

文章提出一种新的并行策略 Elastic Sequence Parallelism，并构建了一个分布式LLM LoongServe以充分释放ESP的潜力。

LoongServe由一组实例及一个 global manager组成。global manager可以动态将这些实例组织成一个个并行组并动态调整各并行组大小，每组并行处理请求。

Global manager 以 iteration 为粒度动态地调度请求、调整 ESP 组和管理分布式 KV Cache 缓存池。

LoongServe Architecture

Elastic Sequence Parallelism(ESP)

Prefill Phase

ESP in Prefill Phase

  1. 划分输入序列给一个并行组的实例
  2. 各实例在本地同时计算得到各自分得token的Q、K、V，计算局部attention并把所得KV存储到本机
  3. 各实例将自己的KV传递给相邻实例，直到各实例的Q与全部KV均进行了attention运算，而后得到本层输出后输入到下一层……

这样做使得模型支持的上下文长度上限由单机提升到多机，同时，由于 KV Cache 块在组内机器上循环传递，为下文 Prefill 阶段的 Scale down 时降低 KV Cache 迁移开销提供机会（做到“zero-cost”）

Decode Phase

在Prefill阶段，并行组的各实例分别存储了部分KV。

  1. 指定组内一个实例作为master，负责计算新生成token的q和kv，保存kv到本机KV Cache中并计算局部attention
  2. master将1.中得到的q发送给所有slave，各slave并行计算局部attention并将结果返回给master
  3. master将结果reduce后做后续的FFN计算……

这种方法只要master有足够的显存存储新生成token的KV就能一直迭代下去。

当master显存不足时，则另外指定组内实例为新的master或对当前并行组做scale up并指定新实例作为master。因此文章中称分布式attention计算方式为 multi-master distributed decoding。

Scale up & Scale down

Lifecycle of requests

组内的计算需求下降时，例如由 prefill 阶段切换到 decode 阶段，可以进行 scale down，降低 DoP（degree of parallelism）。此时，需要将 KV Cache 集中到少数的实例。由于 prefill 阶段的 sequence parallelism 机制，kv 会在组内传递，实例只需要选择性地保存一部分 kv，这部分KV Cache迁移与传递并计算attention重叠，因此说是zero cost。

举个🌰，如下图所示，要将实例 1-3 组成的 并行组scale down为实例 1-2 时，若将 kv1-4 保存在实例 1，kv5-6 保存在实例 2，实例 1 只需要在组间传递 kv 的时候选择保存 kv2，实例 2 只需要选择保存 kv6。

Scale Down

组内的计算或者显存需求上升时，例如随着 decode，可以进行 scale up。此时，需要使新加入的实例可以参与并行计算。由于采用前面提到的 multi-master distributed decoding 机制，只要单个实例有足够的内存来存储新生成的键值张量，就可以通过指定它为 Master 来跨多个实例进行处理。当扩展并行组时，global manager 只需要将新实例添加到 parallel group 中，而不需要迁移现有的KV Cache。

Reference

[^1]: LLM Inference Series: 2. The two-phase process behind LLMs’ responses | by Pierre Lienhart | Medium
[^2]: LLM Inference Series: 3. KV caching explained | by Pierre Lienhart | Medium