图解大模型计算加速系列：vLLM源码解析2，调度器策略(Scheduler)
输入“/”快速插入内容
图解大模型计算加速系列：vLLM源码解析2，调度器策略(Scheduler)
图解大模型计算加速系列：vLLM源码解析2，调度器策略(Scheduler)
用户1329用户13292024年8月20日修改
在上一篇关于vLLM代码整体架构的文章中，我们提到过无论是“离线批处理（同步）”还是“在线流式服务（异步）”，它们都采用了同一个推理内核引擎LLMEngine，其整体架构如下：
其中：
•
在每1个推理阶段中，调度器（Scheduler）决定哪些请求可以参与推理，并为这些请求做好逻辑块->物理块的映射。
•
在每1个推理阶段中，分布式执行者（图中Distributed Workers部分，根据代码，我们将其命名为model_executor会更加合适）接收调度器传来的这些请求，分发到各个worker上去做推理。Worker中的CacheEngine负责实际管理KV Cache；Worker中的model负责加载模型、实行推理，PagedAttention相关的实现和调用就在model下。
这里，每1个推理阶段的定义是：prefill算1个推理阶段，每个decode各算1个推理阶段。在本文中，我们统一用step来表示“1个推理阶段”。
•
在本文中，我们会详细解读调度器（Scheduler）全部细节；
•
在下一篇文章中，我们会详细解读块管理（blockmanager）的全部细节，并以parallel sampling，beam search和prefix caching为例，将上图左半部分全部串一遍
•
在后续文章中，我们会来解读上图右半部分细节（还没来得及拆逻辑，暂时不知道会写几篇）
由于块管理者和调度器在代码上逻辑层层嵌套，所以为了不影响大家对调度器的理解，涉及到块管理者的部分，本文也会给出尽量简明清晰的说明。
一、入口函数
在源码架构篇中我们提过，本系列的介绍思路是：以“离线批处理”作为入口，详细解说内核引擎LLMEngine的各块细节。在此基础上我们再来看“在线流式服务”的运作流程。所以现在，我们先来回顾下离线批处理的调用方式：
代码块
from vllm import LLM, SamplingParams===========================================================================batch prompts===========================================================================prompts = ["Hello, my name is","The president of the United States is","The capital of France is","The future of AI is",]===========================================================================采样参数===========================================================================sampling_params = SamplingParams(temperature=0.8, top_p=0.95)===========================================================================初始化vLLM offline batched inference实例，并加载指定模型===========================================================================llm = LLM(model="facebook/opt-125m")===========================================================================推理===========================================================================outputs = llm.generate(prompts, sampling_params)===========================================================================对每一条prompt，打印其推理结果===========================================================================for output in outputs:prompt = output.promptgenerated_text = output.outputs[0].textprint(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
有两点需要注意：
•
llm = LLM(model="facebook/opt-125m")：实例化了一个离线批处理的vLLM对象。其本质是实例化了一个内核引擎LLMEngine对象。在执行这个步骤时，LLMEngine会执行一次模拟实验（profiling），来判断需要在gpu上预留多少的显存空间给KV Cache block（模拟实验的流程参见源码篇1的3.2节，TODO，大家可以对照着来读源码，本文不再涉及这块源码细节）。
•
推理入口在第24行outputs = llm.generate(prompts, sampling_params)。现在我们进入LLM类下，来看这个generate函数，代码如下：
代码块
vllm/entrypoints/llm.pyclass LLM:"""An LLM for generating texts from given prompts and sampling parameters.
...
"""def init(self,model: str,tokenizer: Optional[str] = None,tokenizer_mode: str = "auto",trust_remote_code: bool = False,tensor_parallel_size: int = 1,dtype: str = "auto",quantization: Optional[str] = None,revision: Optional[str] = None,tokenizer_revision: Optional[str] = None,seed: int = 0,gpu_memory_utilization: float = 0.9,swap_space: int = 4,enforce_eager: bool = False,max_context_len_to_capture: int = 8192,disable_custom_all_reduce: bool = True,**kwargs,) -> None:...==============================================================================使用配置好的engine参数，初始化LLMEngine实例==============================================================================self.llm_engine = LLMEngine.from_engine_args(engine_args, usage_context=UsageContext.LLM_CLASS)==============================================================================用于全局唯一的request_id，在vLLM中内核引擎的处理中，1个prompt视为1个request，分配全局唯一的request_id==============================================================================self.request_counter = Counter()...def generate(self,prompts: Optional[Union[str, List[str]]] = None,
sampling_params: Optional[SamplingParams] = None,prompt_token_ids: Optional[List[List[int]]] = None,
use_tqdm: bool = True,lora_request: Optional[LoRARequest] = None,multi_modal_data: Optional[MultiModalData] = None,) -> List[RequestOutput]:"""Generates the completions for the input prompts.
NOTE: This class automatically batches the given prompts, considering
the memory constraint. For the best performance, put all of your prompts
into a single list and pass it to this method.
Args:
prompts: prompts可以是str，也可以是list[str]
sampling_params: 采样超参，例如温度、top_k等；如果为None则使用vLLM默认的参数
prompt_token_ids: prompt对应的token_id，如果没有提供的话，vllm会调用tokenizer进行 转换
use_tqdm: 是否要展示process bar
lora_request: 如果想请求特定的lora_adapter，可以将它的path等信息包装在该请求中,
但vLLM建议尽量不要使用这种方式，因为私有的lora adapter可能会带来一些
安全性的问题
multi_modal_data: 多模态相关的数据
Returns:
A list of RequestOutput objects containing the generated
completions in the same order as the input prompts.
"""if prompts is None and prompt_token_ids is None:raise ValueError("Either prompts or prompt_token_ids must be ""provided.")if isinstance(prompts, str):Convert a single prompt to a list.prompts = [prompts]if (prompts is not None and prompt_token_ids is not Noneand len(prompts) != len(prompt_token_ids)):raise ValueError("The lengths of prompts and prompt_token_ids ""must be the same.")if sampling_params is None:Use default sampling params.sampling_params = SamplingParams()if multi_modal_data:multi_modal_data.data = multi_modal_data.data.to(torch.float16)============================================================================将request添加到engine中在vLLM内核运算逻辑中，1个prompt算1个request，需要有1个全局唯一的request_id============================================================================num_requests = len(prompts) if prompts is not None else len(prompt_token_ids)for i in range(num_requests):prompt = prompts[i] if prompts is not None else Nonetoken_ids = None if prompt_token_ids is None else prompt_token_ids[i]=======================================================================将每个prompt添加进LLMEngine中，_add_request具体做了以下几件事：- 将每个prompt处理成特定的输入类型（SequenceGroup实例，后文会细说）- 将每个prompt加入Scheduler的waiting队列，等待处理=======================================================================self._add_request(prompt,sampling_params,token_ids,lora_request=lora_request,Get ith image while maintaining the batch dim.multi_modal_data=MultiModalData(type=multi_modal_data.type,data=multi_modal_data.data[i].unsqueeze(0))if multi_modal_data else None,)============================================================================把这个batch的所有prompt都添加完后，执行推理，详情参见_run_engine============================================================================return self._run_engine(use_tqdm)def _add_request(self,prompt: Optional[str],sampling_params: SamplingParams,prompt_token_ids: Optional[List[int]],lora_request: Optional[LoRARequest] = None,multi_modal_data: Optional[MultiModalData] = None,) -> None:每个prompt赋1个request_idrequest_id = str(next(self.request_counter))self.llm_engine.add_request(request_id,prompt,sampling_params,prompt_token_ids,lora_request=lora_request,multi_modal_data=multi_modal_data)def _run_engine(self, use_tqdm: bool) -> List[RequestOutput]: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)===========================================================================如果当前调度器中还有没完成推理的请求（调度器中waiting/running/swapped任一队列非空）===========================================================================outputs: List[RequestOutput] = []while self.llm_engine.has_unfinished_requests():=========================================================================执行1次推理调度（step），决定哪些请求的数据可以参与到这次推理中=========================================================================step_outputs = self.llm_engine.step()for output in step_outputs:=====================================================================如果本step后，有请求已经完成了推理，就将推理结果装进outputs中=====================================================================if output.finished:outputs.append(output)if use_tqdm:pbar.update(1)if use_tqdm:pbar.close()Sort the outputs by request ID.This is necessary because some requests may be finished earlier thanits previous requests.outputs = sorted(outputs, key=lambda x: int(x.request_id))return outputs
总结来说，当我们调用·outputs = llm.generate(prompts, sampling_params)时，它实际做了两件事情：