https://zhuanlan.zhihu.com/p/688551989

在 vllm 中,prefill 和 decoding 不会同时发生。

vllm 的算子经过特殊设计,可以支持不同长度的 batch,不像预训练,batch 作为一个单独的维度,每个 batch 的 seqlen 必须一样。vllm 更像是把多个 batch 拼成一个超长 batch 来推理。

推理流程

准备阶段:维护 waiting 队列和 running 队列。初始时所有请求都在 waiting 队列里。

schedule 阶段:

  1. 遍历 waiting 队列的请求,看当前显存是否放得下。如果放得下,就拿出来,放到 running 队列。

    image.png

  2. 由于此处的请求都是刚刚进入 running 队列,所以进行一次 prefill(而不是 decode)。

image.png

  1. 为执行 decode 做前置检查。尝试为 running 队列中的每个请求,申请一个新的显存块,用来存下一次的 decode 结果。如果不能为每个请求申请显存,就抢占优先级最低的那个。此处为 S4,因为按照入队顺序作为优先级,越早的请求优先级越高。于是,我们把 S4 的 KV cache 时放掉,只保留当前的 token ids,并把 S4 返回 waiting 队列。
  2. 执行 decode。结束后,
  3. 发现有完成的请求:S2 生成了 END token,标志着请求结束。

image.png

释放 S2 的显存。并且,从 waiting 队列中按优先级检查,找到合适的,显存能容纳的请求,比如就是 S4(刚刚步骤3中被回收的那个)。

  1. 此时 S4 没有 KV cache,因此需要单独为 S4 执行一次 prefill。
  2. 把 S4 拼在 S1,S3 后面,然后返回步骤 4:执行 decode。

image.png