vLLM llm_engine.py 源码精读

# 1. 各类 Config 对象直接赋值保存
self.model_config = model_config       # 模型配置
self.cache_config = cache_config       # KV Cache 配置
self.scheduler_config = scheduler_config
self.parallel_config = parallel_config
self.speculative_config = speculative_config
self.decoding_config = decoding_config or DecodingConfig()

# 2. Tokenizer（支持 LoRA 多 tokenizer 组）
self._init_tokenizer()
self.detokenizer = Detokenizer(self.tokenizer)   # L153

# 3. 序列 ID 全局计数器
self.seq_counter = Counter()   # L158

# 4. 分布式执行器（根据设备/并行配置动态选择实现）
self.model_executor = executor_class(...)   # L162

# 5. 初始化 KV Cache（由 executor profiling 决定块数）
self._initialize_kv_caches()   # L174

# 6. 调度器
self.scheduler = Scheduler(scheduler_config, cache_config, lora_config)   # L219

# 7. 指标记录器
self.stat_logger = StatLogger(
    local_interval=_LOCAL_LOGGING_INTERVAL_SEC,   # 5秒
    labels=dict(model_name=model_config.served_model_name),
    max_model_len=self.model_config.max_model_len)   # L223

# 8. 输出处理器（单步 or 多步，取决于 speculative decoding）
self.output_processor = SequenceGroupOutputProcessor.create_output_processor(
    self.scheduler_config, self.detokenizer, self.scheduler,
    self.seq_counter, self.get_tokenizer_for_seq,
    stop_checker=StopChecker(...))   # L231

@nvtx_tag
def step(self) -> List[RequestOutput]:
    # ── 阶段 1：调度 ──────────────────────────────────────────
    seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()
    # scheduler_outputs 包含：
    #   scheduled_seq_groups   本轮参与执行的序列组列表
    #   ignored_seq_groups     超长被忽略的组
    #   blocks_to_swap_in/out  需要在 GPU/CPU 间搬运的 KV block
    #   blocks_to_copy         prefix caching 需要复制的 block
    #   num_lookahead_slots    投机解码的前瞻槽数
    #   running_queue_size     当前 running 队列大小

    # ── 阶段 2：执行（如果有任务）────────────────────────────
    if not scheduler_outputs.is_empty():
        execute_model_req = ExecuteModelRequest(
            seq_group_metadata_list=seq_group_metadata_list,
            blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
            blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
            blocks_to_copy=scheduler_outputs.blocks_to_copy,
            num_lookahead_slots=scheduler_outputs.num_lookahead_slots,
            running_queue_size=scheduler_outputs.running_queue_size,
        )
        output = self.model_executor.execute_model(
            execute_model_req=execute_model_req)
    else:
        output = []   # 队列为空，跳过执行

    # ── 阶段 3：处理输出 ──────────────────────────────────────
    request_outputs = self._process_model_outputs(
        output,
        scheduler_outputs.scheduled_seq_groups,
        scheduler_outputs.ignored_seq_groups,
        seq_group_metadata_list)

    # ── 阶段 4：记录指标 ──────────────────────────────────────
    self.do_log_stats(scheduler_outputs, output)

    return request_outputs

@nvtx_tag
def add_request(
    self,
    request_id: str,
    prompt: Optional[str],
    sampling_params: SamplingParams,
    prompt_token_ids: Optional[List[int]] = None,
    arrival_time: Optional[float] = None,
    lora_request: Optional[LoRARequest] = None,
    multi_modal_data: Optional[MultiModalData] = None,
) -> None:
    # 1. 校验 logprobs 上限
    if arrival_time is None:
        arrival_time = time.time()

    # 2. 编码 prompt → token ids（如果未提供）
    prompt_token_ids = self.encode_request(
        request_id=request_id,
        prompt=prompt,
        prompt_token_ids=prompt_token_ids,
        lora_request=lora_request)

    # 3. 构建 Sequence 对象
    block_size = self.cache_config.block_size
    seq_id = next(self.seq_counter)   # 全局递增的序列 ID
    eos_token_id = self.tokenizer.get_lora_tokenizer(lora_request).eos_token_id
    seq = Sequence(seq_id, prompt, prompt_token_ids, block_size,
                   eos_token_id, lora_request)

    # 4. 克隆 SamplingParams（防止调用方修改影响引擎内部状态）
    sampling_params = sampling_params.clone()
    sampling_params.all_stop_token_ids.add(seq.eos_token_id)
    sampling_params.update_from_generation_config(self.generation_config_fields)

    # 5. 构建 SequenceGroup（一个请求对应一个 SequenceGroup）
    seq_group = SequenceGroup(request_id, [seq], sampling_params,
                              arrival_time, lora_request, multi_modal_data)

    # 6. 提交给调度器的 waiting 队列
    self.scheduler.add_seq_group(seq_group)

@nvtx_tag
def _process_model_outputs(
    self,
    output: List[SamplerOutput],
    scheduled_seq_groups: List[ScheduledSequenceGroup],
    ignored_seq_groups: List[SequenceGroup],
    seq_group_metadata_list: List[SequenceGroupMetadata],
) -> List[RequestOutput]:

    now = time.time()

    # ── 步骤 1：重组输出维度 ──────────────────────────────────
    # 模型输出格式是 [step][seq_group]，重组为 [seq_group][step]
    # 这对于多步执行（speculative decoding）很重要
    output_by_sequence_group = create_output_by_sequence_group(
        sampler_outputs=output,
        num_seq_groups=len(scheduled_seq_groups))

    # ── 步骤 2：逐组处理输出 ──────────────────────────────────
    for scheduled_seq_group, outputs, seq_group_meta in zip(
            scheduled_seq_groups, output_by_sequence_group,
            seq_group_metadata_list):
        seq_group = scheduled_seq_group.seq_group

        # 更新本组已计算的 token 数（用于 chunked prefill 追踪）
        seq_group.update_num_computed_tokens(
            scheduled_seq_group.token_chunk_size)

        # 处理 prompt logprobs（无论是否采样都要更新）
        self.output_processor.process_prompt_logprob(seq_group, outputs)

        # 只有 do_sample=True 时才处理生成的 token
        if seq_group_meta.do_sample:
            self.output_processor.process_outputs(seq_group, outputs)

    # ── 步骤 3：释放已完成的序列组 ───────────────────────────
    self.scheduler.free_finished_seq_groups()

    # ── 步骤 4：构建 RequestOutput ───────────────────────────
    request_outputs: List[RequestOutput] = []
    for scheduled_seq_group in scheduled_seq_groups:
        seq_group = scheduled_seq_group.seq_group
        seq_group.maybe_set_first_token_time(now)   # 记录 TTFT 时间戳
        request_output = RequestOutput.from_seq_group(seq_group)
        request_outputs.append(request_output)

    # 被忽略的组（超过 max_model_len）也需要返回一个错误输出
    for seq_group in ignored_seq_groups:
        request_output = RequestOutput.from_seq_group(seq_group)
        request_outputs.append(request_output)

    return request_outputs

class SequenceGroupOutputProcessor(ABC):
    """处理序列组中新 token id 的接口：
    管理 detokenization、stop checking 以及 scheduler 中的序列 fork/free。

    与 LLMEngine 高度耦合，可视为 LLMEngine 的功能性扩展，
    分离出来是为了简化 LLMEngine 并支持不同解码策略的独立实现。
    """

    @staticmethod
    def create_output_processor(
        scheduler_config: SchedulerConfig,
        detokenizer: Detokenizer,
        scheduler: Scheduler,
        seq_counter: Counter,
        get_tokenizer_for_seq: Callable[[Sequence], PreTrainedTokenizer],
        stop_checker: "StopChecker",
    ):
        # 根据 num_lookahead_slots 决定实现类型
        if scheduler_config.num_lookahead_slots == 0:
            return SingleStepOutputProcessor(...)   # 普通解码
        else:
            return MultiStepOutputProcessor(...)    # 投机解码

    @abstractmethod
    def process_outputs(self, sequence_group: SequenceGroup,
                        outputs: List[SequenceGroupOutput]) -> None:
        """处理新 token id：detokenize、stop check、fork/free 序列。"""

    @abstractmethod
    def process_prompt_logprob(self, seq_group: SequenceGroup,
                               outputs: List[SequenceGroupOutput]) -> None:
        """将 outputs 中的 prompt logprobs 更新到 seq_group。"""

self.output_processor = (
    SequenceGroupOutputProcessor.create_output_processor(
        self.scheduler_config,
        self.detokenizer,
        self.scheduler,
        self.seq_counter,
        self.get_tokenizer_for_seq,
        stop_checker=StopChecker(
            self.scheduler_config.max_model_len,
            self.get_tokenizer_for_seq,
        ),
    ))

def _get_stats(self, scheduler_outputs, model_output=None) -> Stats:
    # ── 类型 1：系统状态（System Stats）──────────────────────
    num_running_sys = len(self.scheduler.running)    # 正在 GPU 上运行的请求数
    num_swapped_sys = len(self.scheduler.swapped)    # 已换出到 CPU 的请求数
    num_waiting_sys = len(self.scheduler.waiting)    # 等待队列中的请求数

    # KV Cache 使用率
    num_free_gpu = self.scheduler.block_manager.get_num_free_gpu_blocks()
    gpu_cache_usage_sys = 1.0 - (num_free_gpu / num_total_gpu)

    # ── 类型 2：迭代级指标（Iteration Stats）────────────────
    # 遍历 scheduled_seq_groups，区分 prefill 组和 decode 组：
    for idx, scheduled_seq_group in enumerate(scheduled_seq_groups):
        group_was_prefill = idx < scheduler_outputs.num_prefill_groups
        if group_was_prefill:
            num_prompt_tokens_iter += scheduled_seq_group.token_chunk_size
            if not seq_group.is_prefill():   # prefill 刚结束 → 记录 TTFT
                latency = seq_group.get_last_latency(now)
                time_to_first_tokens_iter.append(latency)
        else:
            latency = seq_group.get_last_latency(now)  # decode → 记录 TPOT
            time_per_output_tokens_iter.append(latency)

    # ── 类型 3：请求级指标（Request Stats，仅在请求完成时记录）
    if seq_group.is_finished():
        time_e2e_requests.append(now - seq_group.metrics.arrival_time)
        num_prompt_tokens_requests.append(len(seq_group.prompt_token_ids))
        # ... best_of, n, finished_reason 等

    return Stats(now=now, num_running_sys=..., gpu_cache_usage_sys=..., ...)

def log(self, stats: Stats) -> None:
    # 通道 1：每步上报到 Prometheus（实时）
    self._log_prometheus(stats)

    # 累积 token 计数用于吞吐量计算
    self.num_prompt_tokens.append(stats.num_prompt_tokens_iter)
    self.num_generation_tokens.append(stats.num_generation_tokens_iter)

    # 通道 2：每 5 秒上报到 stdout（人工可读摘要）
    if self._local_interval_elapsed(stats.now):
        prompt_throughput = self._get_throughput(self.num_prompt_tokens, stats.now)
        generation_throughput = self._get_throughput(self.num_generation_tokens, stats.now)

        logger.info(
            "Avg prompt throughput: %.1f tokens/s, "
            "Avg generation throughput: %.1f tokens/s, "
            "Running: %d reqs, Swapped: %d reqs, "
            "Pending: %d reqs, GPU KV cache usage: %.1f%%, "
            "CPU KV cache usage: %.1f%%",
            prompt_throughput, generation_throughput,
            stats.num_running_sys, stats.num_swapped_sys,
            stats.num_waiting_sys,
            stats.gpu_cache_usage_sys * 100,
            stats.cpu_cache_usage_sys * 100,
        )
        # 重置累积计数
        self.num_prompt_tokens = []
        self.num_generation_tokens = []
        self.last_local_log = stats.now