vllm源码解析(一)：整体架构与推理代码

最新推荐文章于 2025-03-24 15:10:28 发布

rkmhr_sef

最新推荐文章于 2025-03-24 15:10:28 发布

阅读量1.2k

点赞数 20

分类专栏：面试学习路线阿里巴巴文章标签：架构

本文链接：https://blog.csdn.net/rkmhr_sef/article/details/144870008

版权

vlllm官方代码更新频发,每个版本都有极大变动, 很难说哪个版本好用.
第一次阅读vllm源码是0.4.0版本,对这版圈复杂度极高的调度代码印象深刻
0.4.1对调度逻辑进行重构,完全大变样, 读代码速度快赶不上迭代的速度了。
现在已经更新到0.5.4, 经过长时间观察，发现主要的调度逻辑基本也稳定了下来, 应该可以作为一个固话的版本去阅读。

本文解读依据vllm 0.5.4版本. 没有修改任何代码,大家不必担心夹带私货！
打算以六篇文章的篇幅剖析vllm，希望能对大家有所帮助。

注解代码链接:
https://github.com/yblir/vllm-learn

参考文献：
https://zhuanlan.zhihu.com/p/691038809
https://zhuanlan.zhihu.com/p/681716326

一大模型推理流程

在解析vllm源码前，我们先来回顾下llm推理流程。一个典型的推理过程如下:

在这里插入图片描述

① prefill：预填充阶段，把整段prompt喂给大模型做推理，获得kv-cache并保存。
②decode：大模型本质是个自回归模型，因此生成阶段，首先根据prompt中最后一个token的kv（input token 4）计算获得第一个推理结果（北），并保存对应的kv-cache(output token 1), 这个过程算一次推理；之后将北字作为输入（首次推理的输入是prompt，以后模型输入都是上次的生成token, 当然过程中要用到之前保存的kv-cache），做同样的推理生成京字，直到推理结束。

由于Decode阶段是逐一生成token，因此不能像prefill阶段那样能做大段prompt的并行计算，所以在LLM推理过程中，Decode阶段的耗时一般是更大的，单步生成token的耗时约占总推理时长的90%。

上述推理过程使用到了kv-cache技术，这里有些问题需要解决：
· 随着生成token的增多，kv-cache长度也变大，对gpu显存造成压力
· 生成的token长度无法预知，因此不能提前预知kv-cache所需的存储空间，给推理工作造成很大不确定性

vllm就是为解决上述问题而生，vllm的核心就是如何优化kv-cache，节省显存提高推理吞吐量。
调用方法也很简单，以下是qwen2 vllm推理代码：

# -*- coding: utf-8 -*-
# @Time    : 2024/8/18 20:14
# @Author  : yblir
# @File    : qwen2_vllm_inference.py
# explain  : 
# =======================================================
import os
import sys

sys.path.append('/mnt/e/PyCharm/insteresting/vllm-0.5.4/')

from vllm_module import LLM, SamplingParams
# from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
model_path = '/mnt/e/PyCharm/PreTrainModel/qwen2_15b_instruct'
# model_path = '/media/xk/D6B8A862B8A8433B/data/qwen2-15b-instruct'

params = {"repetition_penalty": 1.1,
          "temperature"       : 0.7,
          'n'                 : 4,
          "top_p"             : 0.8,
          "top_k"             : 20, }
sample_params = SamplingParams(**params)
llm = LLM(model=model_path,
          dtype='half'
            # dtype='float16'
          # 把模型层均分到n个gpu上, 而不是运行n个完整模型
          # tensor_parallel_size=1
          # gpu利用率最大70%
          # gpu_memory_utilization=0.7,
          )
tokenizer = AutoTokenizer.from_pretrained(model_path, )

# 构造模板
prompt = '介绍下京杭大运河'
messages = [
    {'role': 'system', 'content': '你是一个诗人'},
    {'role': 'user', 'content': prompt}
]

text = tokenizer.apply_chat_template(conversation=messages, tokenize=False, add_generation_prompt=True)

messages2 = [
    {'role': 'system', 'content': '你是一个诗人'},
    {'role': 'user', 'content': 'how far you go'}
]

text2 = tokenizer.apply_chat_template(conversation=messages2, tokenize=False, add_generation_prompt=True)

messages3 = [
    {'role': 'system', 'content': '你是一个诗人'},
    {'role': 'user', 'content': '中国首都城市什么名字'}
]

text3 = tokenizer.apply_chat_template(conversation=messages3, tokenize=False, add_generation_prompt=True)
# print(text)
outputs = llm.generate(
        # 当tokenizer.apply_chat_templat中 tokenize为 False 时激活prompts
        prompts=[text,text2,text3],

        # 当tokenizer.apply_chat_templat中 tokenize为 True 时激活prompt_token_ids,与prompts二选一
        # prompt_token_ids=[text,text2,text3],

        sampling_params=sample_params
)

for output in outputs:
    # prompt = output.prompt
    # print(prompt)
    # print(output)
    # print('------------------------------------------')
    for i,item in enumerate(range(4)):
        print(output.outputs[i].text)
        print(output.outputs[i].token_ids)
    print('------------------------------------------
')

看起来很简单吧，似乎只要2步：只要把模型初始化，再调用generate方法就搞定了。实际上这两步的后面是耦合了调度与模型改造的复杂工程，本文将深度剖析潜藏在背后的源码。

二 vllm 原理分析

vllm管理kv-cache的技术称为PagedAttention，原理类似于虚拟内存分页管理技术。
正常推理流程中，生成的token长度无法预知，因此会最大化分配一块连续显存作为kv-cache的存储空间，可能到推理结束时这些空间大部分都用不到，而且这是为当前prompt分配的，其他prompt不能使用，造成极大浪费。
换个思路，如果把显存切分成多个连续小段是否可以呢！动态分配显存小段与生成的kv-cache 之间的存储关系，这样可以最大限度地使用显存，达到提升限度推理吞吐量的目的。
这种方法称为PagedAttention，主要有3个模块组成：logical kv blocks, block table, physical kv blocks. 原理如下图所示，下面我们来逐一解析。
在这里插入图片描述

logical kv blocks：逻辑表，不实际存储kv-cache,可以理解为C++语言中的指针，prefill和decode生成的kv-cache的"地址指针"存储在logical kv blocks，逻辑表对"指针"的存储是连续的。不过在新版vllm中，logical kv这个东西已经删除了，当然逻辑块只是形式上消失了，实际上它依然隐藏在Sequence类的各个属性中，解释起来比较复杂，我们在以后的代码分析中再详解。
physical kv blocks：可理解为实际存储token的物理显存，vllm中一个块默认为16（可以装16个token的k/v值