前期提要与本期导览

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

Untitled

其中:

这里,每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}")

有两点需要注意:

# 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)时,它实际做了两件事情