Sarathi的chunked prefill论文

decode网络结构：

各部分，在各个Batch Size下，耗时：

上图可以看到：

1. Prefill阶段，batch size为1时，GPU算力已满(saturate)，所以提高batch size不能减少单个request的耗时。

2. Decode阶段，因为是memory-bound，提高batch size几乎可以让单个request的耗时线性减少。

3. 耗时大的部分：2个FFN，preproj，postproj。因为这几个是大矩阵计算。attention计算和other(dropout、nomalization、relu等），耗时小。

4. 它的实验里，batch-size为1时，1个decode的平均耗时，是1个prefill的平均耗时，的200倍。（我的实验里是50~100倍）。

5. decode阶段的attention，没有随着batch size增大而减少。（似乎没有打成batch，是一个一个独立计算的）。

吞吐量和batch-size、sequence-length的关系：

1. prefill阶段，B*L>=512 tokens时，GPU算力饱和，吞吐量达到最高，再上不去了。

2. decode阶段，batch-size小的时候，吞吐量随batch-size线性增长。增长到256~512时，GPU算力接近饱和，也涨不上去了。

3. deocde阶段，seq越长，吞吐量越小，因为attention部分耗时和seq长度有关。

计算密度

arithmetic intensity，the amount of compute per memory read/write。

计算次数，除以，<读+写>显存的次数。用来观察是compute bound还是memory bound。

decode在batch-size=256时，计算密度还比prefill在batch-size=1时要低。

chunked prefill的attention mask

多大chunk size合适：

一般设在1K~4K。大些的模型，256也能把GPU打满。

chunked prefill的开销

1. 拆分chunk导致计算密度下降。

解决：给定模型和卡之后，用接近线上的并发数，来profiling不同的chunk-size下的吞吐量和延迟，选择一个最优的chunk size。

2. 读取KV cache的次数上升。

我的疑问：不拆分chunk的做法，和拆分，读取KV cache的次数不是一样的吗？

chunk-size为什么要是2的幂次方

因为CUDA矩阵乘法会进行分块，每个block处理一块。分块的大小是按照block大小来分的。block大小和GPU硬件warp有关。warp是32。因此，如果矩阵大小不是2的整数幂，就会有较明显的GPU算力浪费。

阶跃效应，128->129，延迟跃升：

prefill阶段的attention都是一个一个独立来跑。decode阶段的attention如何batch到一起？

chunked prefill的优点：

1个prefill chunk的256个tokens和多个decoder tokens，的1次前向计算，只读取1遍模型参数。

评测1

baseline是：耗时/decode tokens

实验组是：chunked prefill的耗时，减去，prefill tokens单独打的batch的耗时，最后除以decode tokens。

为什么batch-size越大加速比越小：因为batch-size越大，baseline的计算密集度也上来了。

为什么seq-length越大加速比越小：因为chunked-prefill只对非attention操作有加速效果。seq-length越大，attention的耗时占比越大，chunked-prefill的加速效果就被打薄了。

评测2

2个模型，2种卡。

为什么A100的加速比降了：因为A100的计算密度更大（牛卡主要牛在计算能力上，访存能力比弱卡强不了多少），因此需要更高的chunk-size才能打满GPU算力。

为什么吞吐量只增加20%~30%：因为chunked-prefill只对decode阶段加速了；prefill阶段没有加速。

我的实验观察到的：吞吐量增加不大。延迟显著减少。

评测3

P:D这个比值，也就是requests数据里input prompt tokens和output tokens的数量比值，对chunked-prefill的优势发挥，有影响。

以上纵轴都是和不使用chunked-prefill的baseline实验结果对比得到的加速比。

现象：

1. P:D很小时，很少的decode tokens能搭上prefill chunk的便车，因此和baseline相比没优势。P:D很大时，prefill阶段占主要耗时，搭上便车的这部分decode tokens带来的加速比，被prefill淹没了，因此和baseline比也没优势。所以，才有了上图中，随着P:D的增大，吞吐量先快速上升，再缓慢下降。 

2. 当prefill的chunk，和decode需要的chunk，完全相等时，吞吐量加速比最大，即P/C==D/(B-1)时（C是chunk-size；B-1是因为batch里有1个prefill request）。例如，C=256, B=18时，P/D=C/(B-1)=256/(18-1)=15。

3. 太小的chunk-size，导致计算密度降低，无法打满GPU算力，吞吐量下降。

评测4：

各个部分的耗时占比。（橘黄是baseline，蓝色是chunked-prefill)

差距主要在ffn上，因为矩阵大。其次是preproj，因为生成Q、K、V三组，矩阵也不小。preproj矩阵太小差距不明显。atten反而变慢。

Chunked-prefill VS. Orca iterative scheduling

Orca iterative scheduling：新request来了后（或多个新request)，和正在decoding的tokens放到同1个batch里，推理1次，完成1个（或多个）prefill和decoding tokens的1个step。

等价于把chunk-prefill的chunk-size设的无穷大。

缺点1：导致decode tokens能搭上prefill便车的很少。很多decode tokens还得自己独立计算。因此吞吐量加速比不好。

缺点2：对decode阶段那些requests的延迟造成增大，较明显。

评测5：

1. attention变慢

原文说，chunked prefill会造成对前面chunks的KV cache读取次数增多。我的理解是，不拆分chunk，读取次数，也是这么多，但是，拆分chunk，可能使FlashAttention这种加速可能发挥不出来。

如图，chunk-size为64时，变慢挺明显的。

2. prefill变慢

这里的变慢，包括了：1. KV cache读取次数变多导致的attention变慢；2. 计算密度降低导致的Linear矩阵乘法变慢。

3. 整体吞吐量增大

decode阶段的变快，超过了，prefill阶段的变慢。所以总体吞吐量是增大的。

2的幂次方现象：320，384，448，性能变差，是因为不是2的幂次方。

Other work（系统优化）

VLLM的Paged Attention解决了预先为batch里所有requests分配最大长度的output KV cache，造成的显存浪费问题。

FlexGen提供了显存优化方案：offload，量化，调度。

《Self-attention does not need o(n 2 ) memory, 2022》 降低了attention的显存占用（似乎不流行）。

FlashAttention系列，通过融合kernel，减少了读写显存。

Facebook的xformers对attention的优化，管用。

FasterTransformer做了很多kernel融合的工作。

Orca的iterative-scheduling，避免了decoding阶段末尾的requests padding对GPU算力的浪费。且对完成的request立即返回，不等整个batch的完成。

FastServe提出了抢占式调度。

Triton和Clipper将serving和推理engine分离开。

Other work（模型优化）

Grouped Query Attention：减少了KV cache占用，允许了更大的batch-size。

量化；（减少显存占用，加速）

MoE：减少被激活的参数。

retentive networks？