记录一个优雅的大模型 prefill 超长上下文部署方案 Chunked Pipeline Parallelism，这个方案是 SGLang 博客 (opens new window) 中的一部分。

这篇文章偏随笔，不做完整介绍，细节请自行展开或结合 AI 食用。

1. Chunked Prefill

Chunked Prefill 就是把 prefill 的序列进行分块，分多次完成推理。

Original: Iter0 [ABCDEFGHIJKLMNOPQRSTUVWXYZ]

Chunked Prefill: Iter0 [ABCDEFGH] Iter1 [IJKLMNOP] Iter2 [QRSTUVWX] Iter3 [YZ]

2. Sequence Parallelism

我们有很多个节点，面对超长上下文推理，该怎么发挥多节点的优势呢？一个容易想到的做法是序列并行 Sequence Parallelism。

将 prefill 的序列进行分块，发给各个节点同时进行推理。

Original: rank0 [ABCDEFGHIJKLMNOPQRSTUVWXYZ]

Sequence Parallelism: rank0 [ABCDEFGH] rank1 [IJKLMNOP] rank2 [QRSTUVWX] rank3 [YZ]

但是我们知道 Attention 要在 Sequence Length 维度进行 reduce 操作，导致 Sequence Parallelism 会有通信代价。

3. Pipeline Parallelism

流水线并行 (Pipeline Parallelism) 是将大模型的 transformers 部署到多个节点上，激活值就会在多个节点上流动。

Original: rank0 [layer0 -> layer1 -> layer2 -> layer3]

PP: rank0 [layer0] -> rank1 [layer1] -> rank2 [layer2] -> rank3 [layer3]

当然在 rank1 推理时，rank0 也不能让它闲着，它可以同时进行下一批数据的推理。这样每个节点就会像工厂流水线一样运作起来。

l = layer, b = batch
rank0 [l0×b0 -> l0×b1 -> l0×b2 -> l0×b3]
rank1 [         l1×b0 -> l1×b1 -> l1×b2 -> l1×b3]
rank2 [                  l2×b0 -> l2×b1 -> l2×b2 -> l2×b3]
rank3 [                           l3×b0 -> l3×b1 -> l3×b2 -> l3×b3]

流水线并行虽然能保证吞吐，但是对延迟没有优化效果。

4. Chunked Pipeline Parallelism

将上面的三个方案结合一下，我们得到了 Chunked Pipeline Parallelism 方案。

像 Pipeline Parallelism 那样将 layers 分散在多个节点上；像 Chunked Prefill 那样对序列分块，作为 Pipeline Parallelism 的多个 batch。

l = layer, b = batch

Original Sequence: [ABCDEFGHIJKLMNOPQRSTUVWXYZ]

Chunked Sequence: b0 [ABCDEFGH] b1 [IJKLMNOP] b2 [QRSTUVWX] b3 [YZ]

rank0 [l0×b0 -> l0×b1 -> l0×b2 -> l0×b3]
rank1 [         l1×b0 -> l1×b1 -> l1×b2 -> l1×b3]
rank2 [                  l2×b0 -> l2×b1 -> l2×b2 -> l2×b3]
rank3 [                           l3×b0 -> l3×b1 -> l3×b2 -> l3×b3]

看起来就是很自然的方案，但是它能利用多节点的性能，同时避免了 Sequence Parallelism 的 Attention 通信开销。

可以说是非常优雅地解决了时延问题。