在上一篇关于vLLM代码整体架构的文章中,我们提到过无论是“离线批处理(同步)”还是“在线流式服务(异步)”,它们都采用了同一个推理内核引擎LLMEngine,其整体架构如下:

其中:
这里,每1个推理阶段的定义是:prefill算1个推理阶段,每个decode各算1个推理阶段。在本文中,我们统一用step来表示“1个推理阶段”。
由于块管理者和调度器在代码上逻辑层层嵌套,所以为了不影响大家对调度器的理解,涉及到块管理者的部分,本文也会给出尽量简明清晰的说明。
在源码架构篇中我们提过,本系列的介绍思路是:以“离线批处理”作为入口,详细解说内核引擎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.prompt
generated_text = output.outputs[0].text
print(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,大家可以对照着来读源码,本文不再涉及这块源码细节)。outputs = llm.generate(prompts, sampling_params)。现在我们进入LLM类下,来看这个generate函数,代码如下:# vllm/entrypoints/llm.py
class 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 None
and 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 None
token_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_id
request_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 than
# its previous requests.
outputs = sorted(outputs, key=lambda x: int(x.request_id))
return outputs
总结来说,当我们调用outputs = llm.generate(prompts, sampling_params)时,它实际做了两件事情:
_add_request:将输入数据传给LLMEngine,它具体做了如下事情: