许多LLM的参数大小超过了单个GPU的容量,因此有必要将它们跨分布式GPU进行分割在以模型并行方式下执行。这就要求内存管理器能够处理分布式内存。vLLM通过支持Megatron-LM风格的张量模型并行策略来在分布式设置中有效工作。这种策略遵循SPMD(单程序多数据)执行计划,其中线性层被分割以执行块矩阵乘法,并且GPU通过allreduce操作不断同步中间结果。具体来说,注意力运算在注意力头维度上被分割,每个SPMD进程处理多头注意力中的一个子集。
我们观察到,即使在模型并行执行的情况下,每个模型分片仍然处理相同的输入Token集,因此需要相同位置的KV缓存。因此,vLLM在调度器内设置了一个单一的KV缓存管理器,如图4所示。 不同的GPU workers 共享管理器以及从逻辑块到物理块的映射。这个共同的映射允许 GPU workers 使用调度器为每个输入请求提供的物理块来执行模型。虽然每个GPU worker 有相同的物理块ID,但worker只存储其相应注意力头的KV缓存的一部分。
在每一步中,调度器首先为每个请求准备包含输入Token ID的消息,以及每个请求的块表。 接下来,调度器将此控制消息广播给GPU workers。然后,GPU工作器开始使用输入Token ID执行模型。在注意力层中,GPU工作器根据控制消息中的块表读取KV缓存。在执行过程中,GPU工作器使用all-reduce通信原语同步中间结果,而不需要调度器的协调。最后,GPU工作器将本次迭代的采样Token发送回调度器。总之,GPU工作器不需要在内存管理上进行同步,它们只需要在每个解码迭代开始时以及步骤输入时接收所有的内存管理信息即可。