2. Orca: Iteration Level Scheduling and Selective Batching

Prefill 阶段的计算复杂度

工业常见 dense LLM 的 hidden dimension 大概是 4K 到 16K 量级

无法 batch 的几种情况

其实卡住 batch 的原因都是在 attention 模块的计算部分

• 两个进程都在 prefill 阶段但是二者张量的形状不同;
  • prefill 阶段二者都要计算完整的 $softmax(\frac{QK^t}{\sqrt{d_k}})V$, 如果要一起计算就需要 $X.shape(-1)$ 和 $X.shape(-2)$ 分别相等，也就是张量形状一致;
• 两个进程都在 decode 阶段但是二者的 token id 不同;
  • decode 阶段计算是最后一个 token 对应的 query 去和前面所有的 K 和 V 做匹配，如果二者的 index 不同，说明其前面的 K 和 V 的行数 (sequence id/维度) 不同，从而导致形状不同
• 两个进程处在不同的 p/d 阶段;
  • 计算量不同，一个是矩阵 times 向量; 一个是矩阵 times 矩阵

Orca 的局限性

Orca 的设计目的是实现 iteration level scheduling 的问题，一次 iteration 就是 transformer 完整的模块吐出一个 token 的粒度，那么在这个粒度内部，如果存在 prefill 和 decode 并行的情况，那么我们会注意到问题:

Prefill 的计算复杂度远远高于 Decode phase, 那么 flatten 阶段，比如 $X\mapsto Q, K, V$ ，prefill 的映射其实是 $O(N^2)$ 复杂度的映射，而 decode 的映射是 $O(N)$ 的复杂度，这里 prefill 往往会比 decode 更慢，那么其实我们是拖慢了 decode 的速度，并且在计算 attention 部分的时候 $softmax(\frac{QK^T}{\sqrt{d_k}})V$ 其实也是 decode 的更快，但是我们计算完成之后，进行 FFN 之前我们仍然要进行一次 merge 操作，也就是 decode 需要等 prefill 做完 attention 这里仍然存在少量的延迟;

所以我们看出了在 iter 内部的木桶效应，那么接下来要解决这个，就需要 Sarathi-Serve 那篇论文的讲解了;

同时我们还会想问在 selective batching 的应用背景: Linear 运算对于整体运算的用时贡献相比于 Attention 部分大不大; 这个对比也在 Sarathi-Serve 中有所体现