04-vLLM推理引擎源码解析

# vllm/core/block.py

@dataclass
class PhysicalTokenBlock:
    """
    物理Block: 实际存储KV cache的GPU内存块

    每个Block存储固定数量token的KV cache
    """
    device: Device           # CPU or GPU
    block_number: int        # 物理Block编号
    block_size: int          # Block大小 (token数)
    ref_count: int = 0       # 引用计数 (用于共享)

    def __post_init__(self):
        # 实际的KV tensor在CacheEngine中分配
        pass


@dataclass
class LogicalTokenBlock:
    """
    逻辑Block: 请求视角的Block

    一个逻辑Block可能映射到一个物理Block
    多个逻辑Block可以共享同一个物理Block (prefix caching)
    """
    block_number: int
    block_size: int
    token_ids: List[int] = field(default_factory=list)
    num_tokens: int = 0

    def is_full(self) -> bool:
        return self.num_tokens == self.block_size

    def append_token(self, token_id: int) -> None:
        assert not self.is_full()
        self.token_ids.append(token_id)
        self.num_tokens += 1

    def get_num_empty_slots(self) -> int:
        return self.block_size - self.num_tokens


class BlockTable:
    """
    Block表: 维护逻辑Block到物理Block的映射

    类似操作系统的页表
    """

    def __init__(self):
        self.logical_to_physical: Dict[int, int] = {}

    def add_mapping(self, logical_block_num: int, physical_block_num: int):
        self.logical_to_physical[logical_block_num] = physical_block_num

    def get_physical_block(self, logical_block_num: int) -> int:
        return self.logical_to_physical[logical_block_num]

    def get_block_table(self) -> List[int]:
        """返回物理Block编号列表 (传给CUDA kernel)"""
        return [self.logical_to_physical[i]
                for i in range(len(self.logical_to_physical))]

# vllm/core/block_manager.py

class BlockSpaceManager:
    """
    Block空间管理器

    管理GPU和CPU上的物理Block分配和释放
    """

    def __init__(
        self,
        block_size: int,
        num_gpu_blocks: int,
        num_cpu_blocks: int,
        watermark: float = 0.01,
        sliding_window: Optional[int] = None,
    ):
        self.block_size = block_size
        self.num_gpu_blocks = num_gpu_blocks
        self.num_cpu_blocks = num_cpu_blocks

        # GPU Block池: 空闲的物理Block
        self.gpu_allocator = BlockAllocator(
            Device.GPU, block_size, num_gpu_blocks
        )

        # CPU Block池 (用于swap)
        self.cpu_allocator = BlockAllocator(
            Device.CPU, block_size, num_cpu_blocks
        )

        # 每个sequence的Block表
        self.block_tables: Dict[int, BlockTable] = {}

        # watermark: 预留的最小空闲Block数
        self.watermark_blocks = int(num_gpu_blocks * watermark)

    def can_allocate(self, seq_group: SequenceGroup) -> AllocStatus:
        """检查是否能为新请求分配Block"""
        # 计算需要的Block数
        num_required_blocks = self._get_num_required_blocks(seq_group)

        # 检查空闲Block是否足够
        num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()

        if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:
            return AllocStatus.OK
        elif num_free_gpu_blocks >= num_required_blocks:
            return AllocStatus.LATER  # 等待其他请求完成
        else:
            return AllocStatus.NEVER  # 请求太长,永远无法分配

    def allocate(self, seq_group: SequenceGroup) -> None:
        """为新请求分配Block"""
        seq = seq_group.get_seqs()[0]

        # 计算需要的Block数
        num_prompt_blocks = (len(seq.prompt_token_ids) + self.block_size - 1) // self.block_size

        # 分配物理Block
        block_table = BlockTable()

        for logical_idx in range(num_prompt_blocks):
            # 分配物理Block
            physical_block = self.gpu_allocator.allocate()
            block_table.add_mapping(logical_idx, physical_block.block_number)

        self.block_tables[seq.seq_id] = block_table

    def can_append_slot(self, seq_group: SequenceGroup) -> bool:
        """检查是否能追加一个新token"""
        seq = seq_group.get_seqs()[0]
        block_table = self.block_tables[seq.seq_id]

        # 检查最后一个Block是否有空间
        last_block_idx = len(block_table.logical_to_physical) - 1
        last_block = self.gpu_allocator.get_block(
            block_table.get_physical_block(last_block_idx)
        )

        if last_block.get_num_empty_slots() > 0:
            return True

        # 需要分配新Block
        return self.gpu_allocator.get_num_free_blocks() >= 1

    def append_slot(self, seq: Sequence) -> Optional[Tuple[int, int]]:
        """追加一个slot用于新生成的token"""
        block_table = self.block_tables[seq.seq_id]

        # 检查最后一个Block
        if len(block_table.logical_to_physical) > 0:
            last_logical_idx = len(block_table.logical_to_physical) - 1
            last_physical_idx = block_table.get_physical_block(last_logical_idx)
            last_block = self.gpu_allocator.get_block(last_physical_idx)

            if last_block.get_num_empty_slots() > 0:
                # 最后一个Block还有空间
                return None

        # 需要分配新Block
        physical_block = self.gpu_allocator.allocate()
        new_logical_idx = len(block_table.logical_to_physical)
        block_table.add_mapping(new_logical_idx, physical_block.block_number)

        return (new_logical_idx, physical_block.block_number)

    def free(self, seq: Sequence) -> None:
        """释放sequence的所有Block"""
        block_table = self.block_tables.pop(seq.seq_id)

        for physical_block_num in block_table.logical_to_physical.values():
            self.gpu_allocator.free(physical_block_num)

    def swap_out(self, seq_group: SequenceGroup) -> Dict[int, int]:
        """将GPU Block换出到CPU"""
        mapping = {}

        for seq in seq_group.get_seqs():
            block_table = self.block_tables[seq.seq_id]

            for logical_idx, gpu_block_num in block_table.logical_to_physical.items():
                # 分配CPU Block
                cpu_block = self.cpu_allocator.allocate()

                # 记录映射 (用于数据拷贝)
                mapping[gpu_block_num] = cpu_block.block_number

                # 更新Block表
                block_table.add_mapping(logical_idx, cpu_block.block_number)

                # 释放GPU Block
                self.gpu_allocator.free(gpu_block_num)

        return mapping

    def swap_in(self, seq_group: SequenceGroup) -> Dict[int, int]:
        """将CPU Block换入到GPU"""
        mapping = {}

        for seq in seq_group.get_seqs():
            block_table = self.block_tables[seq.seq_id]

            for logical_idx, cpu_block_num in block_table.logical_to_physical.items():
                # 分配GPU Block
                gpu_block = self.gpu_allocator.allocate()

                # 记录映射
                mapping[cpu_block_num] = gpu_block.block_number

                # 更新Block表
                block_table.add_mapping(logical_idx, gpu_block.block_number)

                # 释放CPU Block
                self.cpu_allocator.free(cpu_block_num)

        return mapping

# vllm/core/scheduler.py

class Scheduler:
    """
    请求调度器

    决定:
    1. 哪些请求可以在当前step执行
    2. 哪些请求需要preempt (换出)
    3. 哪些请求可以swap回来
    """

    def __init__(
        self,
        scheduler_config: SchedulerConfig,
        cache_config: CacheConfig,
    ):
        self.scheduler_config = scheduler_config
        self.cache_config = cache_config

        # Block管理器
        self.block_manager = BlockSpaceManager(
            block_size=cache_config.block_size,
            num_gpu_blocks=cache_config.num_gpu_blocks,
            num_cpu_blocks=cache_config.num_cpu_blocks,
        )

        # 请求队列
        self.waiting: Deque[SequenceGroup] = deque()  # 等待处理
        self.running: Deque[SequenceGroup] = deque()  # 正在运行
        self.swapped: Deque[SequenceGroup] = deque()  # 已换出到CPU

    def add_seq_group(self, seq_group: SequenceGroup) -> None:
        """添加新请求到等待队列"""
        self.waiting.append(seq_group)

    def schedule(self) -> Tuple[SchedulerOutputs, bool]:
        """
        主调度函数

        返回:
        - SchedulerOutputs: 调度结果
        - bool: 是否有prompt (prefill)
        """
        # 1. 首先处理swapped队列 (优先恢复之前被抢占的请求)
        self._schedule_swapped()

        # 2. 处理running队列 (decode阶段的请求)
        self._schedule_running()

        # 3. 处理waiting队列 (新请求,prefill阶段)
        self._schedule_waiting()

        # 构建输出
        return self._build_scheduler_outputs()

    def _schedule_waiting(self) -> None:
        """调度等待队列中的请求"""
        while self.waiting:
            seq_group = self.waiting[0]

            # 检查是否能分配Block
            alloc_status = self.block_manager.can_allocate(seq_group)

            if alloc_status == AllocStatus.OK:
                # 可以分配,移到running队列
                self.waiting.popleft()
                self._allocate(seq_group)
                self.running.append(seq_group)
            elif alloc_status == AllocStatus.LATER:
                # 暂时不能分配,等待
                break
            else:  # NEVER
                # 请求太长,无法处理
                self.waiting.popleft()
                self._abort_seq_group(seq_group)

    def _schedule_running(self) -> None:
        """调度正在运行的请求"""
        # 检查是否有足够的Block继续生成
        preempted = []

        for seq_group in self.running:
            # 检查是否能追加新token
            if not self.block_manager.can_append_slot(seq_group):
                # Block不足,需要抢占
                preempted.append(seq_group)

        # 抢占优先级最低的请求
        for seq_group in reversed(preempted):
            self._preempt(seq_group)
            self.running.remove(seq_group)
            self.swapped.append(seq_group)

    def _schedule_swapped(self) -> None:
        """调度被换出的请求"""
        while self.swapped:
            seq_group = self.swapped[0]

            # 检查是否能swap回来
            if self.block_manager.can_swap_in(seq_group):
                self.swapped.popleft()
                self._swap_in(seq_group)
                self.running.append(seq_group)
            else:
                break

    def _preempt(self, seq_group: SequenceGroup) -> None:
        """抢占请求 (swap out到CPU)"""
        preemption_mode = self.scheduler_config.preemption_mode

        if preemption_mode == PreemptionMode.SWAP:
            # 换出到CPU
            self._swap_out(seq_group)
        else:
            # 重新计算 (丢弃KV cache,之后重算)
            self._recompute(seq_group)

    def _allocate(self, seq_group: SequenceGroup) -> None:
        """为请求分配Block"""
        self.block_manager.allocate(seq_group)

    def _swap_out(self, seq_group: SequenceGroup) -> Dict[int, int]:
        """换出到CPU"""
        return self.block_manager.swap_out(seq_group)

    def _swap_in(self, seq_group: SequenceGroup) -> Dict[int, int]:
        """从CPU换入"""
        return self.block_manager.swap_in(seq_group)

# vllm/attention/backends/flash_attn.py

class FlashAttentionBackend(AttentionBackend):
    """
    使用FlashAttention的Attention后端
    """

    @staticmethod
    def get_impl_cls() -> Type["FlashAttentionImpl"]:
        return FlashAttentionImpl


class FlashAttentionImpl(AttentionImpl):

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: FlashAttentionMetadata,
        k_scale: float = 1.0,
        v_scale: float = 1.0,
    ) -> torch.Tensor:
        """
        执行PagedAttention

        Args:
            query: [num_tokens, num_heads, head_size]
            key: [num_tokens, num_kv_heads, head_size]
            value: [num_tokens, num_kv_heads, head_size]
            kv_cache: [2, num_blocks, block_size, num_kv_heads, head_size]
            attn_metadata: 包含block_tables, context_lens等
        """
        num_tokens, num_heads, head_size = query.shape

        # 更新KV Cache
        if kv_cache is not None:
            key_cache = kv_cache[0]
            value_cache = kv_cache[1]

            # 将新的KV写入cache
            ops.reshape_and_cache(
                key,
                value,
                key_cache,
                value_cache,
                attn_metadata.slot_mapping,
            )

        # Prefill (prompt处理)
        if attn_metadata.prefill_metadata is not None:
            output = flash_attn_varlen_func(
                q=query,
                k=key,
                v=value,
                cu_seqlens_q=attn_metadata.prefill_metadata.seq_start_loc,
                cu_seqlens_k=attn_metadata.prefill_metadata.seq_start_loc,
                max_seqlen_q=attn_metadata.prefill_metadata.max_seq_len,
                max_seqlen_k=attn_metadata.prefill_metadata.max_seq_len,
                softmax_scale=self.scale,
                causal=True,
            )

        # Decode (生成阶段)
        if attn_metadata.decode_metadata is not None:
            output = paged_attention_v1(
                out=output,
                query=query,
                key_cache=key_cache,
                value_cache=value_cache,
                num_kv_heads=self.num_kv_heads,
                scale=self.scale,
                block_tables=attn_metadata.decode_metadata.block_tables,
                context_lens=attn_metadata.decode_metadata.context_lens,
                block_size=self.block_size,
                max_context_len=attn_metadata.decode_metadata.max_context_len,
            )

        return output

# vllm/engine/llm_engine.py

class LLMEngine:
    """
    vLLM核心引擎

    管理模型、调度器、Worker
    """

    def __init__(
        self,
        model_config: ModelConfig,
        cache_config: CacheConfig,
        parallel_config: ParallelConfig,
        scheduler_config: SchedulerConfig,
        device_config: DeviceConfig,
        lora_config: Optional[LoRAConfig] = None,
    ):
        self.model_config = model_config
        self.cache_config = cache_config

        # 初始化调度器
        self.scheduler = Scheduler(scheduler_config, cache_config)

        # 初始化Worker (可能是分布式)
        self._init_workers(
            model_config,
            parallel_config,
            device_config,
        )

        # 初始化KV Cache
        self._init_cache()

    def _init_workers(self, model_config, parallel_config, device_config):
        """初始化Worker进程"""
        if parallel_config.world_size == 1:
            # 单GPU
            self.workers = [Worker(model_config, parallel_config, 0)]
        else:
            # 多GPU (使用Ray分布式)
            self.workers = []
            for rank in range(parallel_config.world_size):
                worker = ray.remote(Worker).remote(
                    model_config, parallel_config, rank
                )
                self.workers.append(worker)

    def _init_cache(self):
        """Profile并初始化KV Cache"""
        # Profile: 确定能分配多少Block
        num_gpu_blocks, num_cpu_blocks = self._profile_num_available_blocks()

        self.cache_config.num_gpu_blocks = num_gpu_blocks
        self.cache_config.num_cpu_blocks = num_cpu_blocks

        # 初始化Cache
        self._run_workers("init_cache_engine", cache_config=self.cache_config)

    def add_request(
        self,
        request_id: str,
        prompt: str,
        sampling_params: SamplingParams,
    ) -> None:
        """添加新请求"""
        # Tokenize
        prompt_token_ids = self.tokenizer.encode(prompt)

        # 创建Sequence
        seq = Sequence(seq_id, prompt_token_ids)
        seq_group = SequenceGroup(request_id, [seq], sampling_params)

        # 添加到调度器
        self.scheduler.add_seq_group(seq_group)

    def step(self) -> List[RequestOutput]:
        """
        执行一步推理

        这是vLLM的核心循环
        """
        # 1. 调度: 决定哪些请求参与本次推理
        scheduler_outputs = self.scheduler.schedule()

        if scheduler_outputs.is_empty():
            return []

        # 2. 准备输入
        execute_model_req = self._prepare_model_input(scheduler_outputs)

        # 3. 执行模型
        output = self._run_workers(
            "execute_model",
            execute_model_req=execute_model_req,
        )

        # 4. 处理输出
        return self._process_model_outputs(output, scheduler_outputs)

    def _prepare_model_input(
        self,
        scheduler_outputs: SchedulerOutputs
    ) -> ExecuteModelRequest:
        """准备模型输入"""
        # 收集所有要处理的序列
        seq_group_metadata_list = []

        for seq_group in scheduler_outputs.scheduled_seq_groups:
            seq_data = {}
            for seq in seq_group.get_seqs():
                seq_data[seq.seq_id] = SequenceData(
                    prompt_token_ids=seq.prompt_token_ids,
                    output_token_ids=seq.output_token_ids,
                )

            # Block表
            block_tables = {}
            for seq in seq_group.get_seqs():
                block_table = self.scheduler.block_manager.get_block_table(seq)
                block_tables[seq.seq_id] = block_table

            metadata = SequenceGroupMetadata(
                request_id=seq_group.request_id,
                is_prompt=scheduler_outputs.prompt_run,
                seq_data=seq_data,
                sampling_params=seq_group.sampling_params,
                block_tables=block_tables,
            )
            seq_group_metadata_list.append(metadata)

        return 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,
        )

    def _process_model_outputs(
        self,
        output: SamplerOutput,
        scheduler_outputs: SchedulerOutputs
    ) -> List[RequestOutput]:
        """处理模型输出"""
        request_outputs = []

        for seq_group, outputs in zip(
            scheduler_outputs.scheduled_seq_groups,
            output.outputs
        ):
            for seq, output in zip(seq_group.get_seqs(), outputs):
                # 获取采样结果
                token_id = output.samples[0].output_token
                logprob = output.samples[0].logprobs

                # 追加到序列
                seq.append_token_id(token_id, logprob)

                # 分配新的slot (如果需要)
                self.scheduler.block_manager.append_slot(seq)

                # 检查是否完成
                if self._check_stop(seq, seq_group.sampling_params):
                    seq.status = SequenceStatus.FINISHED_STOPPED

            # 构建输出
            request_output = RequestOutput(
                request_id=seq_group.request_id,
                prompt=seq_group.prompt,
                outputs=[...],
                finished=seq_group.is_finished(),
            )
            request_outputs.append(request_output)

        return request_outputs