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

石小疯 · 4 天前

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值），图中展示每个block大小为4。每个block内部是一连的，但block之间是不一连的，那么如何才能与logical保持对齐呢？这就需要block table了
block table：存储logical与physical关系的映射表。如logical block0 -> physical block7. block table除了记载映射关系，还记载当前block槽位填充情况。如physical block7已经填满，因此filled==4； physical block1 槽位填充了3个，再填入一个father单词的token，filled会变为4。

多batch并行推理时，会有logical blocks映射到同一个physical blocks上，大家看图就能理解:

我们考虑别的一个题目，llm推理偶然会有多个输出，这种情况PagedAttention该如何操作呢？
多输出有两种情况：
Parallel Sampling: 如果指定了n个输出，就把prompt复制n份，拼成一个batch喂给模型做推理。这时会产生prompt 的kv-cache重复存储，对这个重复的优化是别的题目，这里不展开了。
Beam Search：集束搜索，每个decode阶段，产生top k个token（k也被称为束宽），对应着当前时候的top k个序列。它们的前置token也会有大量的kv-cache重复。

前面提到Parallel Sampling模式会把prompt复制 n份，如Figure 8所示，对应sample A1和sample A2，它们各自维护一套自己的logical blocks，由于内容完全类似，它们共享一套physical blocks，此中每个physical block对应引用计数ref count都为2. 进入推理阶段后，A1和A2各自独立做推理，如果生成了类似的token，会把新token kv-cache参加共享的physical block中。如果生成了差别的token（如图中的mothers和fathers），会触发copy-on-write机制，即在gpu上开辟一个新的block，如physical block1复制内容到block2，之后各自再装入生成的差别token, 同时，block1 计数-1，block2计数+1。这种操作符合vllm核心思想：节流KV cache显存，对于类似数据对应的KV cache，能复用则尽量复用；无法复用时，再考虑开辟新的物理空间。
vllm也有对Beam Search有优化，但这不是本文重点，暂时忽略。
目前为止，我们仅回首与源码解析相关知识，PagedAttention还有很多东西没有讲到，说太多技术点反而会让人迷失在各种细节中，有兴趣可以自行去查资料相识。
三 vllm中一些根本概念

在解析源码前，我们还需要清楚理解vllm一些概念的意思。
3.1 vllm 数据结构

如第一章中图片上展示的，一个prompt的典型例子如下：

<|im_start|>system
你是一个诗人<|im_end|>
<|im_start|>user
中国首都是<|im_end|>
<|im_start|>assistant

复制代码

正常推理过程中，1个请求（batchsize）可能包罗多个prompts，在vllm中，一个prompt才被看做一个请求；一个prompt可能输出多个outputs，这时每对prompt -> output序列称为一个seq序列，每条seq都维护着独立的status,可理解为当前时候所处的推理状态，推理是否完结等:

WAITING：正在waiting队列中。waiting队列中的序列都没有做过prefill。
RUNNING：正在running队列中，即已经开始做推理。
SWAPPED：正在swapped队列中，表示此时gpu资源不足，相关的seq_group被抢占。

固然还有finished状态，重要记载因何种缘故起因导致finished：

FINISHED_STOPPED：正常实行完毕，例如碰到符号，该seq的推理正常结束了
FINISHED_LENGTH_CAPPED：由于seq的长度达到最大长度限制，而结束推理
FINISHED_ABORTED：因不正常状态，而被制止的推理。例如客户端断开毗连，则服务器会制止相关seq的推理
FINISHED_IGNORED：因prompt过长而被制止实行的推理。本质上也是受到长度限制

从上面可以看出单独管理seq有点复杂，所以我们需要统一管理一个prompt和它对应的所有outputs，称为一个seq_group。seq_group是vllm推理管理数据的根本单元。vllm 中设定一个seq_group中所有seq共享共一个prompt。这些变量的包罗关系如下：

这张图是从某位大佬的文章中取来的，他用的版本应该是0.4.0，在0.5.4版本中，Sequence类已删除了logical_token_blocks属性和_append_tokens_to_blocks方法，对应的功能转移到其他代码中了。不过从图中，仍能清晰看到二者的从属关系。
3.2 调度原则

前面提到vllm的核心是对kv-cache的优化，而这种优化是通过调度系统Scheduler来完成的。
Scheduler维护着三个双端队列, 在3.1也有提到：
waiting，running，swapped。
每完成一次推理，都要对这三个队列进行动态调整，在下一次推理时实现最大限度提升吞吐量的目标.
这三个队列的作用如下：

waiting: 所有输入的模型的prompt都会被参加waiting队列中，这是输入数据的入口。这时的seq只有一条，就是prompt，连prefill都还没做，不管在外部设置了多输出还是Beam Search，此时只有这一条数据。waiting队列另一个数据来源是running队列。
running: 存储着上一次被送去做推理seq_groups，在下一次做推理前，要对running队列中的seq_groups做检查，看系统是否有充足资源让它们留在队列中继续做下一次推理。如果当前系统资源不满足做一次推理，就把seq_group一条条pop()出来，转移到waiting或swapped队列中，直到满足下一次推理的资源需求。固然，running队列也会从waiting和swapped队列拿数据过来做推理，至于转移，怎么拿，接下来我们会详细分析。
swapped: 可以理解为失败者集散地，都是不满足条件（gpu blocks资源不足或某些推理参数超过阈值），被从running队列中踢出去的，等条件满足时，还会被从新参加running队列做推理。

上面提到了三个队列间的交互，它们的交互依据就是vllm的调度原则。
vllm调度的原则可以总结如下，相识了下面的处理处罚逻辑，就能理解三个队列的用途及它们之间相互转移数据的规则：

先来的请求先被服务（First-Come-First-Serve, FCFS）
如有抢占的需要，后来的请求先被抢占（preemption）

FCFS大家都很好理解，我们来看下对preemption的处理处罚：
抢占发生在推理阶段，vllm核心是最大限度优化吞吐量，推理过程中gpu显存不足或推理的tokens和seqs数目超过设定阈值，都会发生抢占，即暂时制止一些任务的实行，开释gpu上与它们相关的kv-cache,等资源充足，再规复它们的实行。针对preemption有两种处理处罚方式：
① 如果parallel sampling=1，直接开释所有physical blocks，将任务重新放回wait队列（放到队列头部，下一次开始取它），重新从prefill阶段开始做推理。
② 如果parallel sampling>1, 如果将它们直接丢掉，那未免过于浪费, 先把处理处罚的好的blocks交换到CPU上，等gpu显存充足，再把这些blocks从CPU加载回来。
上面提到超过推理参数超过阈值也会导致抢占，这里的阈值指每次答应推理的最大seqs数目和最大tokens数目。在vllm 0.5.4版本中，由一个类budget的对象管理，每次推理前都要重新构建一个budget类对象，统计seqs和tokens数目是否越界。
至此，我们可以得出结论，判断一个seq_group是否被抢占的因素有三个：gpu blocks数目是否充足，当前调度能处理处罚的seqs和tokens是否超过数目阈值。
3.3 vllm推理流程

有了上面的知识储备，我们就能理解一个seq从开始到结束，整个生命周期内的历程：

① 每条prompt处理处罚成Sequence对象，然后Sequence包装成seq_group，这条seq_group会存入waiting队列。此时只有一条seq，就是prompt，连预填充prefill都没做。status为waiting
② 调度器选中这条seq_group做推理，图中我们展示两种情况，4输出和1输出，因此会产生4条seq和1条seq, 此中4 seq共享prompt，status为running。
③ 推理一段时间后，gpu blocks资源不足或tokens或seqs数目超出阈值，发生抢占现象。多输出的seq_group相关kv blocks会被swap out到CPU上；而单输出的seq_group则会把相关的（prefill和decode）kv blocks开释，将seq_group重新放回waiting队列，就像什么都没发生过，幸运的是会被放到waiting队列最前面。
④ 系统资源充足了，被swapped的seq_group会从CPU上swap_in gpu继续推理；单输出seq_group则从prefill开始重新奋斗。
⑤⑥ 多输出肯定会出现某条seq先推理结束，此时还活跃的seq数减1，变为3个，当某条seq推理完结，会被标记为finished, 以后不再调用资源处理处罚它，只有等seq_group中所有seq都推理结束，该seq_group才算推理完成。

留意：并不是每个seq_group都会经历抢占，如果系统资源充足，会跳过抢占的，此时的实行序次为：①②⑤⑥ 或 ①②⑤
四推理代码解析

经过漫长的知识铺垫，我们终于能看到vllm的核心代码了
我们先略过初始化阶段，直接从generate开始，遇到重要知识点再现场分析

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
)

复制代码

vllm推理代码真的很简洁，只要一行代码就行，背后逻辑就极其复杂了，而且每个版本都有很大改动。

vllm/entrypoints/llm.py
1. def generate(
2. self,
3. prompts: Union[Union[PromptInputs, Sequence[PromptInputs]], Optional[Union[str, List[str]]]] = None,
4. sampling_params: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
5. prompt_token_ids: Optional[Union[List[int], List[List[int]]]] = None,
6. use_tqdm: bool = True,
7. lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
8. prompt_adapter_request: Optional[PromptAdapterRequest] = None,
9. guided_options_request: Optional[Union[LLMGuidedOptions, GuidedDecodingRequest]] = None
10. ) -> List[RequestOutput]:
12. if self.llm_engine.model_config.embedding_mode:
13. raise ValueError(
14. "LLM.generate() is only supported for generation models (XForCausalLM)."
15. )
17. # cast 表面看是类型转换,但实践上没做任何事,直接原样返回, 那么这个调用的意义是什么 ?
18. if prompt_token_ids is not None:
19. inputs = self._convert_v1_inputs(
20. prompts=cast(Optional[Union[str, List[str]]], prompts),
21. prompt_token_ids=prompt_token_ids,
22. )
23. else:
24. inputs = cast(Union[PromptInputs, Sequence[PromptInputs]], prompts)
26. if isinstance(guided_options_request, dict):
27. if len(guided_options_request) > 1:
28. raise ValueError(
29. "You can only use one guided decoding but multiple is "
30. f"specified: {guided_options_request}")
31. guided_options_request = GuidedDecodingRequest(**guided_options_request)
33. if sampling_params is None:
34. # Use default sampling params.
35. sampling_params = SamplingParams()
37. # 校验入参，并将一个batchsize中的每条prompt处理成Sequence对象，然后Sequence包装成SequenceGroup组，
38. # 1. prompt->seq->seq_group, 2. 将seq_group加入合适gpu维护的scheduler的waiting队列,等待处理
39. self._validate_and_add_requests(
40. inputs=inputs,
41. params=sampling_params,
42. lora_request=lora_request,
43. prompt_adapter_request=prompt_adapter_request,
44. guided_options=guided_options_request)
46. # 首先从scheduler的waiting队列取数据，加入到running队列，再从running队列
47. # 中取数据推理，若物理blocks不够用，从running转入swap队列
48. outputs = self._run_engine(use_tqdm=use_tqdm)
50. return LLMEngine.validate_outputs(outputs, RequestOutput)
复制代码

generate 代码中最重要的模块有两个，_validate_and_add_requests（数据预处理处罚），_run_engine（实际推理）。
4.1 数据预处理处罚

vllm/entrypoints/llm.py
1. def _validate_and_add_requests(
2. self,
3. inputs: Union[PromptInputs, Sequence[PromptInputs]],
4. params: Union[SamplingParams, Sequence[SamplingParams], PoolingParams, Sequence[PoolingParams]],
5. lora_request: Optional[Union[Sequence[LoRARequest], LoRARequest]],
6. prompt_adapter_request: Optional[PromptAdapterRequest],
7. guided_options: Optional[GuidedDecodingRequest] = None,
8. ) -> None:
9. # 如果输入是一条prompt,而不是list,会在此处自动转换为list
10. if isinstance(inputs, (str, dict)):
11. # Convert a single prompt to a list.
12. inputs = [inputs]
14. ...
16. # Add requests to the engine.
17. # 遍历每一条prompt，1个prompt算1个request，需要有1个全局唯一的request_id
18. for i, request_inputs in enumerate(inputs):
19. self._add_request(
20. request_inputs,
21. params[i] if isinstance(params, Sequence) else params,
22. lora_request=lora_request[i] if isinstance(lora_request, Sequence) else lora_request,
23. prompt_adapter_request=prompt_adapter_request)
25. # ==========================================================================================================
26. def _add_request(
27. self,
28. inputs: PromptInputs,
29. params: Union[SamplingParams, PoolingParams],
30. lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
31. prompt_adapter_request: Optional[PromptAdapterRequest] = None
32. ) -> None:
33. # 每个prompt赋1个全局唯一的request_id
34. request_id = str(next(self.request_counter))
35. self.llm_engine.add_request(
36. request_id,
37. inputs,
38. params,
39. lora_request=lora_request,
40. prompt_adapter_request=prompt_adapter_request)
复制代码
vllm/engine/llm_engine.py
==========================================================================================================
1. def add_request(
2. self,
3. request_id: str, # 每个请求的唯一id,在vLLM内部，1条prompt算1个请求，会附给1个请求id
4. inputs: PromptInputs, # prompt
5. params: Union[SamplingParams, PoolingParams], # 用于采样的参数（温度、topk等）
6. arrival_time: Optional[float] = None, # 请求到达的时间。如果是None，则用当前系统时间
7. lora_request: Optional[LoRARequest] = None, # 如果是用lora模型做推理，相关的lora请求
8. trace_headers: Optional[Mapping[str, str]] = None,
9. prompt_adapter_request: Optional[PromptAdapterRequest] = None,
11. ) -> None:
12. if lora_request is not None and not self.lora_config:
13. raise ValueError(f"Got lora_request {lora_request} but LoRA is not enabled!")
14. # 设置该请求的到达时间
15. if arrival_time is None:
16. arrival_time = time.time()
18. # processed_inputs:dict,= {'prompts':xxx,'prompts_token_ids':xxx,'multi_modal_data':None}
19. processed_inputs = self.process_model_inputs(
20. request_id=request_id,
21. inputs=inputs,
22. lora_request=lora_request,
23. prompt_adapter_request=prompt_adapter_request)
24. # 1. prompt->seq->seq_group, 2. 将seq_group加入合适gpu维护的scheduler的waiting队列,等待处理
25. self._add_processed_request(
26. request_id=request_id,
27. processed_inputs=processed_inputs,
28. params=params,
29. arrival_time=arrival_time,
30. lora_request=lora_request,
31. prompt_adapter_request=prompt_adapter_request,
32. trace_headers=trace_headers,
33. )
复制代码

我们看到，经过疯狂套娃后，最终实际干活的是self.process_model_inputs和self._add_processed_request这两个方法，其他全都是中心商~

vllm/engine/llm_engine.py
1. def process_model_inputs(
2. self,
3. request_id: str,
4. inputs: PromptInputs,
5. lora_request: Optional[LoRARequest] = None,
6. prompt_adapter_request: Optional[PromptAdapterRequest] = None,
7. ) -> LLMInputs:
8. if isinstance(inputs, str):
9. inputs = {"prompt": inputs}
11. if "prompt_token_ids" not in inputs:
12. # 这个函数就是为了拿到self.tokenizer
13. tokenizer = self.get_tokenizer_group("prompts must be None if skip_tokenizer_init is True")
14. # 文字prompt编码成token_id
15. prompt_token_ids = tokenizer.encode(request_id=request_id,
16. prompt=inputs["prompt"],
17. lora_request=lora_request)
18. # 如果入参前已经做好token_ids,直接取出来用
19. else:
20. prompt_token_ids = inputs["prompt_token_ids"]
22. # 使用未合并的lora才会走进入这个判断分支
23. if prompt_adapter_request:
24. prompt_token_ids = [0] * prompt_adapter_request.prompt_adapter_num_virtual_tokens + prompt_token_ids
26. # LLMInputs继承自TypedDict,将入参转换为字典
27. # llm_inputs = {'prompts':xxx,'prompts_token_ids':xxx,'multi_modal_data':None}
28. llm_inputs = LLMInputs(prompt_token_ids=prompt_token_ids,
29. prompt=inputs.get("prompt"),
30. multi_modal_data=inputs.get("multi_modal_data"))
32. # todo 使用functools.partial高阶用法,返回的是一个固定llm_inputs参数的函数，真的好用吗？
33. # 目前觉得这个函数会拖慢速度,因为每个prompt都要经过这里获得模型架构的操作
34. # 目前这个函数,经过多层调用后, 最后原样返回,没对llm_inputs做任何操作
35. return self.input_processor(llm_inputs)
复制代码

process_model_inputs功能是把输入的prompt转换为token_id, 效果以字典形式输出：
llm_inputs = {‘prompts’:xxx,‘prompts_token_ids’:xxx,‘multi_modal_data’:None}

# ==========================================================================================================
def _add_processed_request(
self,
request_id: str,
processed_inputs: LLMInputs,
params: Union[SamplingParams, PoolingParams],
arrival_time: float,
lora_request: Optional[LoRARequest],
prompt_adapter_request: Optional[PromptAdapterRequest],
trace_headers: Optional[Mapping[str, str]] = None,
) -> None:
# Create the sequences.
block_size = self.cache_config.block_size
# self.seq_counter是在类中初始化，所以可以为每条seq生成不重复的id，
# seq_id与request_id是两个独立的变量
seq_id = next(self.seq_counter)
eos_token_id = self._get_eos_token_id(lora_request)
# seq 包含当前prompt的各种信息:token_id,status(waiting,...), 占用blocks数量(逻辑,物理数量相同)
seq = Sequence(seq_id, processed_inputs, block_size, eos_token_id,
lora_request, prompt_adapter_request)
# 将seq和采样参数合并为seq_group
# --------------------------------------------------------------------------------------------------------------
if isinstance(params, SamplingParams):
seq_group = self._create_sequence_group_with_sampling(
request_id,
seq,
params,
arrival_time=arrival_time,
lora_request=lora_request,
trace_headers=trace_headers,
prompt_adapter_request=prompt_adapter_request)
elif isinstance(params, PoolingParams):
seq_group = self._create_sequence_group_with_pooling(
request_id,
seq,
params,
arrival_time=arrival_time,
lora_request=lora_request,
prompt_adapter_request=prompt_adapter_request)
else:
raise ValueError("Either SamplingParams or PoolingParams must be provided.")
# --------------------------------------------------------------------------------------------------------------
# Add the sequence group to the scheduler with least unfinished seqs.
# 获得当前每个gpu上还没推理结束的seq_group数量: len(self.waiting) + len(self.running) + len(self.swapped)
costs = [scheduler.get_num_unfinished_seq_groups() for scheduler in self.scheduler]
# 找出工作量最少的调度器
min_cost_scheduler = self.scheduler[costs.index(min(costs))]
# 将当前seq_group加入这个调度器中(根据self.scheduler初始过程可知,每个gpu维护一个调度器,
# 这条代码的意思就是当前seq_group由工作量最少的gpu负责推理)
min_cost_scheduler.add_seq_group(seq_group)

复制代码

用户输入的prompt经过_validate_and_add_requests处理处罚后，会封装为seq_group，然后将seq_group参加合适gpu维护的scheduler的waiting队列, 等待处理处罚。
seq和seq_group是vllm推理的根本单元，是两个非常重要的概念，我们来直观感受下它们的数据格式：
seq:

seq_group:
初始阶段seqs_dict都只有一个元素，经过prefill后，会扩展为n个元素（n是输出outputs数目）

4.2 推理流程

def _run_engine(
self, *, use_tqdm: bool
) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
# Initialize tqdm.
if use_tqdm:
num_requests = self.llm_engine.get_num_unfinished_requests()
pbar = tqdm(
total=num_requests,
desc="Processed prompts",
dynamic_ncols=True,
postfix=f"est. speed input: {0:.2f} toks/s, output: {0:.2f} toks/s",
)
# Run the engine.
outputs: List[Union[RequestOutput, EmbeddingRequestOutput]] = []
total_in_toks = 0
total_out_toks = 0
# 如果当前调度器中还有没完成推理的请求（调度器中waiting/running/swapped任一队列非空）
while self.llm_engine.has_unfinished_requests():
# 执行1次推理调度（step），决定哪些请求的数据可以参与到这次推理中，step输出本次推理结果
step_outputs = self.llm_engine.step()
# 一次step推理后，如果有请求已经完成了推理，将推理结果装进outputs中，
for output in step_outputs:
if output.finished:
outputs.append(output)
if use_tqdm:
if isinstance(output, RequestOutput):
# Calculate tokens only for RequestOutput
total_in_toks += len(output.prompt_token_ids)
in_spd = total_in_toks / pbar.format_dict["elapsed"]
total_out_toks += sum(len(stp.token_ids) for stp in output.outputs)
out_spd = total_out_toks / pbar.format_dict["elapsed"]
pbar.postfix = (
f"est. speed input: {in_spd:.2f} toks/s, output: {out_spd:.2f} toks/s"
)
pbar.update(1)
if use_tqdm:
pbar.close()
# Sort the outputs by request ID.
# This is necessary because some requests may be finished earlier than
# its previous requests.
return sorted(outputs, key=lambda x: int(x.request_id))

复制代码

整个推理engine中，最重要的是self.llm_engine.step()，封装了所有的调度，推理和后处理处罚代码。

def step(self) -> List[Union[RequestOutput, EmbeddingRequestOutput]]:
# 多GPU并行推理时走AsyncLLMEngine分支。如果进入当前LLMEngine,性能会下降，这里会抛出异常。
if self.parallel_config.pipeline_parallel_size > 1:
raise NotImplementedError(
"Pipeline parallelism is only supported through AsyncLLMEngine "
"as performance will be severely degraded otherwise.")
# 上述if判断表明，只有一个GPU可用。因此self.scheduler也只有一个元素，是当前GPU的调度
# 该函数调用改变调度的内部状态(self.running、self.swapped 和 self.waiting)
seq_group_metadata_list, scheduler_outputs = self.scheduler[0].schedule()
...
return request_outputs

复制代码

step中利用的调度代码如下：

def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]:
# Schedule sequence groups.
# This function call changes the internal states of the scheduler
# such as self.running, self.swapped, and self.waiting.
# 该函数调用改变调度的内部状态(self.running、self.swapped 和 self.waiting）
scheduler_outputs = self._schedule()
...
return seq_group_metadata_list, scheduler_outputs

复制代码

调度系统是vllm代码的核心，接下来，我们花单独一篇文章，详细解读self._schedule()的内部逻辑。

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。

		自动登录	找回密码
密码			立即注册

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

本帖子中包含更多资源

0 个回复

快速回复

楼主热帖

标签云

浏览过的版块