今天读了一篇论文：Fast Distributed Inference Serving for Large Language Models，来自 PKU 的 Xin Jin 组（也是一个很厉害的系统 / MLSys 组了）。论文的 arxiv 链接：https://arxiv.org/pdf/2305.05920.pdf ，主要内容是讲用进程管理里的经典方式 MLFQ 来管理 LLM Serving 里的句子。
Abstract & Introduction
传统的 LLM serving systems 采用 run-to-completion 的方式来处理 inference jobs，这有着两个大问题：
- head-of-line blocking。一个 large job（即 output length 很长的 job）将会运行很长时间，以至于 block 了后续的 short jobs。
- long JCT（job completion time）
- 在此之前，Orca 被认为是 sota，它首先采用了 iteration-level scheduling（在每次 iteration 结束时候都 schedule 一次，然后重新改动 current processing batch）。然而，它采用先来先服务（FCFS）策略来处理 inference jobs。
这篇文章则提出了一个新的名为 FastServe 的分布式 inference serving system for LLMs，其最大特征使用一个 skip-join Multi-Level Feedback Queue scheduler 来实现一个抢占式调度器。
- MLFQ 是一个很经典的减小平均 JCT 的手法，每个 job 先进入最高级别队列，随着其时间片的使用逐渐放低到低优先级队列。
- 而由于 LLM inference 的一些特点（即其分为 prefill 阶段和 decode 阶段，prefill 阶段生成第一个 token 用时相对较长。在长 prompt 短 output 时 prefill 时间甚至是 dominate 的），这里特地加上了一个 skip-join。相比传统的 MLFQ，加上 skip-join 后新 job 不会直接进入最高优先级队列，而是根据 prefill time 来加入到合适的优先级。
- 这种抢占式调度会带来额外的 memory overhead——每个 unfinished job 的 KV Cache 都得暂时保存在内存中。如果采用简单的方式（cache 满了，暂时不添加新的 job）可以解决问题，但是又会带来 head-of-line blocking；这篇文章设计了一个较高效的 GPU 内存管理机制，它会主动把低优先级 jobs 的 KV Cache 放到 host memory（CPU）中，并且采用了 pipelining 和 asynchronous 到 memory 操作来提高效率。
- FastServe 还采用了 tensor parallelism 和 pipeline parallelism。KV Cache manager 会将 KV Cache 分割到多个 GPU 上（distributed），并且在 GPU 和 host 之间 swapping。
Scheduling
这里先来介绍几种 scheduing 策略：
- FCFS（First-Come-First-Serve）：先来先服务，也即 run-to-completion。最朴素的策略，没有 preemption 机制，也就是一个 job 一旦进入 running 状态就得等它跑完。这会有严重的 head-of-line blocking 问题。
- Naive MLFQ：朴素的 MLFQ，最开始所有 job 都放在最高优先级队列，所以 J1、J2、J3 的 T1 都会先被执行。
- Skip-Join MLFQ：也即本文提出的方法。由于有着 skip-join 的能力，最开始也会做一次类似 shortest remaining time 的 scheduling 操作（也即不是先跑 J1，而是先跑 J2、J3）。注意到 J2、J3 的 T2 比 J1 还要优先。
- SRPT（Shortest Remaining Processing Time）：最理想的 scheduling 方式，也就是永远优先执行当前剩余时间最少的 job。但是在实际情况我们是不知道任务的剩余执行时间的，因此需要用 MLFQ 对其进行近似。
Skip-Join MLFQ 算法:
- Input: 各级队列 Q1, …, Qk，新到达的 jobs J_in，preempted 的 jobs J_pre
- Output: 要被执行的任务 J_out
- Process newly arrival jobs: 将 J_in 的 jobs 根据 next tier time 插到 Qi 中
- Process preempted jobs: 对于那些已经被 preempted 的 jobs，我们直接输出新生成的 token（作为对比，没有 preempted 机制的 scheduler 必须等一个 job 跑完才能输出所有 tokens）。然后查看是否需要 demotion（即下放到更低优先级的队列）。
- 由于这种机制也会出现 starving 现象，我们需要一个 promote 机制来定时地把一些 jobs promote 上去。详细来讲，每个 job 会设置一个 starve timer，如果它在 waiting 状态等待超过了某个 threshold，就直接把它提到最高级队列 Q1。系统管理者可以调整这个 threshold 来在 performance 和 starvation 之间做一个取舍。
- 最后，scheduler 会（按队列优先级从高到低）选择不超过 max batch size 个 jobs 进行执行。
Proactive Key-Value Cache Management
Problem:
Preemption 增加了 system 中 ongoing jobs 的数量。
假设记一个 LLM inference job 的 input length 为 s，output length 为 t，hidden dim 为 h，transformer 层数为 l，并且所有的权重都是 FP16，那么这个 job 占据的 KV Cache 空间大小就为 $4 \times lh(s+t)$。如果用 GPT-3 175B 来计算，就需要 2.3GB 的 memory。
FCFS 有一个好处，就是 waiting 中的 jobs 一定都没有 prefill，也就是说真正占用 KV Cache 的永远是当前正在处理的那些 jobs。根据测量，skip-join MLFQ 的内存开销可能有 FCFS 的 7x 大小。
以下是作者提出的一些 Strawman solutions：
- 如果 GPU 不够，直接推迟新到达的 jobs 的执行。这个方法会使得 MLFQ 退化成 FCFS。
- kill low-priority jobs。其实也就是 recomputation。这个做法有两个问题：第一，重算把中间状态丢了，需要重新 prefill，这又需要额外的计算资源和时间；第二，这可能会带来 deadlocks。当一个 low-priority job 被 kill 后，由于 promotion 机制，被 kill 的 job 可能会被提到最高优先级被执行，这时候它又会去 kill 之前那个 kill 它的 job。
作者提出，一个 key observation 是只有当相应的 job 被 schedule 时，其对应的 KV cache 才需要在 GPU 中。FastServe 可以将 inactive job 对应的 KV cache offload 到 host memory 上，然后upload 必要的 KV cache 到 GPU 上。对 A100 上的 GPT-3 175B 来说，decoding phase 生成一个 token 大概需要 250ms，而在 host 和 GPU 中 transfer KV Cache 大概是 36ms（打满 PCIe 4.0 * 16。这里没说是一个 token 还是一个句子）。
Swapping Strategy:
FastServe 采用一个 proactive（主动） 的 swapping 策略。作为对比，被动的策略是当 KV cache 满了之后才 swap；而主动的策略则是保留一些 idle 的 KV cache 为新到达的 jobs 做准备。当新的 job 到达后，它可以立刻拿到 slot 而不需要付出 offloading 的代价。
Job 的 offloading 和 uploading 是基于 ENST（estimated nest scheduled time）来决定的。很自然地，ENST 越大的 job 会被优先 offload。
这里有一个问题，还是之前提到的 promotion 机制。它意味着低优先级的 job 也可能被先执行。于是 FastServe 将 promote 的时间也引入了 ENST 的考虑之中。
形式化地，对于 job i，它考虑两方面因素：promote job i 的时间，以及所有比 job i 执行优先的 jobs 的执行时间。一个优先于 i 的 job j 的执行时间可以记作：
其中 $quantum(k)$ 是优先级为 k 的队列大小。那么比 i 优先的所有 job 执行时间求和（有点奇怪，这里为什么要求和？而且这里好像也不是很明确“时刻”和“时间”的概念）就得到
最后再和 promote 时间取一个 min