vllm代码快速上手

原文：https://zhuanlan.zhihu.com/p/6462326972

0 背景

最近在做long context的工作，其中一个重点是尝试打破内存墙 。需要对vllm代码做比较大的修改，但团队都是一堆算法，于是大家一起从头学习vllm。
看了一堆vllm的代码解读文章，有点收获，但不多，都不能快速建立一个vllm的认知框架，让我能快速上手做修改，看的我很暴躁。

直到后面翻到下面的文章，才清晰不少。

• HelloWorld的vllm系列文章：https://zhuanlan.zhihu.com/p/681716326

• vllm的meetup ppt：https://docs.google.com/presentation/d/1QL-XPFXiFpDBh86DbEegFXBXFXjix4v032GhShbKf3s/edit#slide=id.g288188ac1bf_2_270

1 LLMEngine

LLMEngine就是vllm的核心 class，从vllm运行代码一直讲到LLMEngine

1.1 vllm运行demo

from vllm import LLM, SamplingParams
if __name__ == '__main__':
 model_path = "/data/root/LLM/Qwen1.5-14B-Chat"
 model = LLM(model=model_path, tensor_parallel_size=1, trust_remote_code=True, max_model_len=10000, enforce_eager=True, 
 gpu_memory_utilization=0.5, block_size=32)
 sampling_params= SamplingParams(temperature=0, max_tokens=1, prompt_logprobs=20)
 
 prompt = "今天天气怎么样?"
 response = model.generate(prompt, sampling_params, use_tqdm=False)[0]
 print(response, '\n\n', response.outputs)

LLM是入口class

1.2 LLM class

代码路径——vllm/entrypoints/llm.py

(这里不放github路径，大家自己找对应版本下载吧，换版本后代码改动都比较大）

1.2.1 初始化
核心在于LLMEngine的初始化

1.2.2 执行
提供了几种接口做不同的生成，generate，beam_search，chat，encode。
这些函数核心都是调用LLMEngine做事情，核心代码是self._run_engine(use_tqdm=use_tqdm)。

_run_engine的核心代码就是 self.llm_engine.step()，这个LLMEngine就是vllm的核心类了。

1.3 LLMEngine class

代码路径——vllm/engine/llm_engine.py

1.3.1 初始化

vllm meetup的ppt，写的非常清晰了，这里直接贴出来。

a] Initialize & load model

b] Profile memory usage

c] Pre-allocate KV Blocks

1.3.2 step()

llm_engine执行的核心函数

 def step(self) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
 """Performs one decoding iteration and returns newly generated results.

 .. figure:: https://i.imgur.com/sv2HssD.png
 :alt: Overview of the step function
 :align: center

 Overview of the step function.

 Details:
 - Step 1: Schedules the sequences to be executed in the next
 iteration and the token blocks to be swapped in/out/copy.

 - Depending on the scheduling policy,
 sequences may be `preempted/reordered`.
 - A Sequence Group (SG) refer to a group of sequences
 that are generated from the same prompt.

 - Step 2: Calls the distributed executor to execute the model.
 - Step 3: Processes the model output. This mainly includes:

 - Decodes the relevant outputs.
 - Updates the scheduled sequence groups with model outputs
 based on its `sampling parameters` (`use_beam_search` or not).
 - Frees the finished sequence groups.

 - Finally, it creates and returns the newly generated results.
"""

注释写的很清晰，分成三步。用更简介的话来表达，就三个事情，预处理，模型执行，后处理。

2 预处理

2.1 llama1的预处理

多个句子，长度不一样，肯定不能直接丢给模型做推理。llama1的做法，找到最短的句子，然后并行执行 。
那有的token重复计算了怎么办？
跳过（我当时看的时候笑了半天，但换我，我估计也这么写，简单）

llama1的generate代码：https://github.com/meta-llama/llama/blob/ef351e9cd9496c579bf9f2bb036ef11bdc5ca3d2/llama/generation.py#L186

        for cur_pos in range(min_prompt_len, total_len):
            logits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
            if temperature > 0:
                probs = torch.softmax(logits[:, -1] / temperature, dim=-1)
                next_token = sample_top_p(probs, top_p)
            else:
                next_token = torch.argmax(logits[:, -1], dim=-1)

            next_token = next_token.reshape(-1)
            # only replace token if prompt has already been generated
            next_token = torch.where(
                input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token
            )
            tokens[:, cur_pos] = next_token
            if logprobs:
                token_logprobs[:, prev_pos + 1 : cur_pos + 1] = -F.cross_entropy(
                    input=logits.transpose(1, 2),
                    target=tokens[:, prev_pos + 1 : cur_pos + 1],
                    reduction="none",
                    ignore_index=pad_id,
                )
            eos_reached |= (~input_text_mask[:, cur_pos]) & (
                next_token == self.tokenizer.eos_id
            )
            prev_pos = cur_pos
            if all(eos_reached):
                break

此刻，你要是个经验丰富的MLSyser，肯定就觉得这个做法也太蠢了，我要优化它。于是，你去调研业务场景的输入都是怎么样的。

2.2 真实世界的输入

有这么三个特点
a] 长度不同，甚至差距很大
b] 请求有波峰
To C的服务，调用大头都是白天，不同产品的波峰还不一样。
办公类产品，用户是上班用。角色扮演类产品，用户肯定是下班后跟自己的二次元老婆聊天。
c] kv cache 总是不够用 机器太贵，kv cache就这么一点，但模型越来越大，输入长度越来越长。

2.3 vllm的预处理

参考文章：https://zhuanlan.zhihu.com/p/681716326 https:/1716326
基于上面三个需求，vllm做了如下的方式来优化，核心代码路径——vllm/core/scheduler.py

2.3.1 排队队列(scheduler)

MLSyser大兄弟肯定很熟悉，业务有波峰很正常，波峰很高也很正常。
但加机器要钱，支持的并发就这么些，你全打给机器，机器就死给你看。你要把这些撑不住的请求丢了，业务方（例如我）一定来干你。
好消息是，内存没那么贵，那么就搞个队列，把处理不完的请求放到队列里（内存），去排队。

Scheduler中有 3 个队列
waiting（接受到的新请求会先放入 waiting 队列）
running（被调度的请求）
swapped 队列（swapped 队列用于存放被抢占的请求，即当请求处于生成阶段时，但由于空间的不足，需暂时将 running 队列中优先级低的请求移到 swapped 队列）。
Scheduler 会按照先到先处理（first come first served）的原则从 waiting 队列中选择请求放入 running 队列

不同业务的排队最优解不同，本科的计算机系统，一堆这样的设计方案 （这是MLSyser大兄弟的领域，我就不深入展开了）

代码如下，可以看到三个队列都是deque

 # Create the block space manager.
 self.block_manager = BlockSpaceManagerImpl(
 block_size=self.cache_config.block_size,
 num_gpu_blocks=num_gpu_blocks,
 num_cpu_blocks=num_cpu_blocks,
 sliding_window=self.cache_config.sliding_window,
 enable_caching=self.cache_config.enable_prefix_caching)

 # Sequence groups in the WAITING state.
 # Contain new prefill or preempted requests.
 self.waiting: Deque[SequenceGroup] = deque()
 # Sequence groups in the RUNNING state.
 # Contain decode requests.
 self.running: Deque[SequenceGroup] = deque()
 # Sequence groups in the SWAPPED state.
 # Contain decode requests that are swapped out.
 self.swapped: Deque[SequenceGroup] = deque()
 # Sequence groups finished requests ids since last step iteration.
 # It lets the model know that any state associated with these requests
 # can and must be released after the current step.
 # This is used to evict the finished requests from the Mamba cache.
 self._finished_requests_ids: List[str] = list()

2.3.2 虚拟内存 （BlockSpaceManager)

从代码可以看到，waiting，running，swapped三个队列上面还有一个 BlockSpaceManagerImpl对象，这就是vllm中的虚拟内存。 kv cache的具体优化靠page attention，放到worker模块下面，后面会提到。

大学的操作系统课程里有个虚拟内存的东西。
电脑的内存也是有限的，并且因为内存分配 的问题，实际的物理内存分配往往是分散的。我们写代码的时候，去做这种零散内存的操作，是危险且麻烦的。
虚拟内存就做了封装，让我们像操作一个大的、连续的内存空间一样编写程序，但实际的物理内存可能是分散的并且大小有限。
除此之外，还有很多好处，这里都是当年的考试题，就不展开了。

正如vllm meetup ppt里面写的，manage KV cache like OS virtual memory

实际的调用方式

2.3.3 selective batch

如上图，还是vllm的meetup ppt，vllm每计算一次token，都会盘一下当前要怎么重新分配资源，根据kv cache的情况，最大化利用计算资源。
而不是像llama1那样，一旦开跑，就要等最长的跑完才行。有点像多线程的生产者和消费者队列进阶版。

具体代码的解析，参看这篇文章：https://zhuanlan.zhihu.com/p/692276562 2276562
流程图画的非常清晰。

事实上这样也有性能上的优化空间
一个batch中做prefill和做decode的请求有多少条是不确定的，只是大体按照先来后到的原则做动态组装。这就造成了一些问题：
如果一个batch中做prefill 的请求非常多，或者做prefill的请求非常长，那么prefill tokens会占据大量计算资源，使得整个batch变成compute-bound 。
如果一个batch中做decode 的请求非常多（比如当所有的请求都没做完推理时，或者请求队列 中没有新序列可以调度时），这个batch就可能变成memory-bound的。随机的batch同样可能产生pp并行气泡。

所以，引入了chunked prefill 来做优化，具体可以参考这篇文章，这里不做展开。 https:// https://mp.weixin.qq.com/s/_nm2Fwz2FlkcLuXnWDB9PA Fwz2FlkcLuXnWDB9PA

3 模型执行

核心在于两点，kv cache怎么优化，model怎么更好的load和infer。个人觉得，最重要的就是kv cache，这里的挑战实在是太大了

3.1 kv cache的挑战

最近两年内，kv cache的优化，会是横亘在几乎所有大模型公司面前的一个挑战。这里从flash attention的视角来展开，它认为kv cache优化分成两种。

[a] FLOPS（Floating Point Operations Per Second）
达到同样效果的情况下，降低计算量。如各种efficient transformer 的优化。

[b]MAC（memory access cost）
gpu是三层级的内存结构

可以看到SRAM 存储量级要远小于HBM，但读取速度远高于HBM。这种性能的差异，导致我们如果频繁的做SRAM和HBM之间的内存交互，会极大的降低效率。
根据小群讨论，SRAM的价格也没有那么贵，核心卡点是带宽慢，intel的pcie 太慢了，老黄看不下去自己撸的nvlink 。

flash attention优化的就是这块，传统transformer有大量的hbm 读写。FlashAttention 将参与计算的矩阵进行分块送进SRAM，来提高整体读写速度（减少了HBM读写）。
但transformer的softmax是没有分块计算的能力，依赖一个全局的分母项。Flash attention ，从3-pass变成了1-pass，来实现了分块计算的能力。
这块就不展开了，vllm直接调用的flash attention ，大家可以自己看细节，参考下面这两篇文章。
https://zhuanlan.zhihu.com/p/642962397 2962397
https://