vLLM：高性能大语言模型推理框架源码解析与最佳实践

最新推荐文章于 2025-09-28 09:48:35 发布

原创

最新推荐文章于 2025-09-28 09:48:35 发布 · 2k 阅读

15 ·

CC 4.0 BY-SA版权

文章标签：

#语言模型 #人工智能 #自然语言处理

该文章已生成可运行项目，

vLLM：高性能大语言模型推理框架源码解析与最佳实践

引言
快速上手
2.1. 安装配置
2.2. 基本用法
核心调用流程分析
3.1. 总体调用链路概述
3.2. 核心组件与类层次结构
3.3. 初始化阶段详细流程
3.4. 推理阶段详细流程
3.5. 完整调用链路示例
3.6. 关键调用路径总结
vLLM 关键工作机制
4.1. PagedAttention 机制
4.2. 连续批处理技术
4.3. CUDA 图捕获与加速
4.4. KV 缓存管理
4.5. 内存优化策略
4.6. 并行计算优化
架构与类关系图
5.1. 整体架构概览
5.2. 核心类结构图
5.3. 关键组件交互流程
5.4. 并行与分布式架构
5.5. 内存管理架构
高级 Python 语法应用
6.1. 类型注解与泛型编程
6.2. 动态类解析与反射机制
6.3. 方法委托与魔术方法
6.4. 弱引用与特殊调用模式
6.5. 类方法与工厂模式
6.6. 装饰器高级应用
6.7. 上下文管理与多重上下文
性能优化与最佳实践
7.1. 示例应用：构建高性能API服务
7.2. 性能优化策略
7.3. 部署最佳实践
7.4. 常见问题与解决方案
总结与展望
8.1. 技术创新点总结
8.2. 与其他框架对比
8.3. 未来发展趋势
8.4. 实践建议

附录

A. 关键源码目录结构
B. 参考资源

1. 引言

vLLM 是一个高性能的大语言模型推理框架，通过创新的 PagedAttention 机制和优化的内存管理，实现了高吞吐量和低延迟的文本生成服务。本文将深入分析 vLLM 的源码，重点关注其调用流程、高级 Python 语法应用以及核心工作机制，帮助读者更好地理解和使用这一强大工具。

2. 快速上手

2.1 安装配置

vLLM 支持多种硬件平台，安装方式也略有不同：

# 基本安装（CUDA 支持）
pip install vllm

# 从源码安装
git clone https://github.com/vllm-project/vllm.git
cd vllm
pip install -e .

2.2 基本用法

下面是一个简单的 vLLM 使用示例，展示了如何加载模型并生成文本：

# 导入必要的库
from vllm import LLM, SamplingParams

# 初始化模型
llm = LLM(model="/opt/data/models/Llama-3.2-1B-Instruct")

# 设置生成参数
sampling_params = SamplingParams(
    max_tokens=100,         # 最大生成长度
    temperature=0.7,        # 控制随机性
    top_p=0.9,              # 控制取样范围
    stop=["</s>", "\n\n"]   # 停止标记
)

# 准备输入提示
prompts = [
    "介绍一下北京的旅游景点",
    "写一首关于春天的诗",
]

# 生成文本
outputs = llm.generate(prompts, sampling_params)

# 打印结果
for output in outputs:
    print(f"Prompt: {
     
     output.prompt}")
    print(f"Generated: {
     
     output.outputs[0].text}")
    print("-" * 50)

3. 核心调用流程分析

3.1 总体调用链路概述

vLLM 框架的调用流程可以分为初始化阶段和推理阶段两个主要部分。整个流程涉及多个核心组件的协同工作，包括 LLM Engine、Worker、Scheduler、ModelRunner 等类。以下是端到端完整调用链路的详细分析：

用户 API 调用
  ↓
LLMEngine 初始化
  ↓ 
模型加载与配置初始化
  ↓
Worker 池初始化
  ↓
请求处理与调度
  ↓
Tokenizer 处理输入
  ↓
ModelRunner 执行推理
  ↓
KV 缓存管理
  ↓
结果获取与后处理
  ↓
返回生成结果

3.2 核心组件与类层次结构

vLLM 的架构由以下核心组件组成：

LLMEngine：整个框架的中枢，负责协调各组件工作
Worker：实际执行模型计算的组件
ModelRunner：负责模型的前向传播计算
Scheduler：请求调度器，决定何时处理哪些请求
SamplingParams：控制文本生成的参数
Request/Response：请求和响应的数据结构
SequenceGroup/Sequence：表示生成序列的数据结构
BlockManager：管理 KV 缓存的物理块
Tokenizer：负责文本标记化

类层次结构图：

LLMEngine
  ├── AsyncEngineDeadlockMonitor
  ├── ModelConfig
  ├── DeviceConfig
  ├── ParallelConfig
  ├── Scheduler
  │    ├── SchedulerConfig
  │    └── SchedulingStrategy
  ├── RequestTracker
  └── WorkerPool
       ├── Worker
       │    ├── ModelRunner
       │    │    ├── CUDAGraph
       │    │    └── AttentionState
       │    └── BlockSpaceManager
       │         └── BlockAllocator
       └── ModelReplica

3.3 初始化阶段详细流程

3.3.1 LLMEngine 初始化

LLMEngine 是整个 vLLM 的核心，负责协调各组件工作。它的初始化流程如下：

def __init__(
    self,
    model_config: ModelConfig,
    scheduler_config: SchedulerConfig,
    device_config: Optional[DeviceConfig] = None,
    parallel_config: Optional[ParallelConfig] = None,
    **kwargs,
) -> None:
    # 1. 设置基本配置
    self.model_config = model_config
    self.scheduler_config = scheduler_config
    if device_config is None:
        device_config = DeviceConfig()
    self.device_config = device_config
    if parallel_config is None:
        parallel_config = ParallelConfig()
    self.parallel_config = parallel_config
    
    # 2. 初始化策略配置
    self.num_nodes = parallel_config.world_size // parallel_config.tensor_parallel_size
    self.distributed_init_method = kwargs.get("distributed_init_method", None)
    self.max_logprobs = kwargs.get("max_logprobs", None)
    
    # 3. 初始化组件
    self.tokenizer = self._init_tokenizer()
    self.schedulers = self._init_schedulers()
    self.request_tracker = RequestTracker()
    
    # 4. 初始化 Worker 池
    self.worker_pool = WorkerPool(
        model_config=model_config,
        parallel_config=parallel_config,
        scheduler_config=scheduler_config,
        device_config=device_config,
        local_rank=parallel_config.local_rank,
    )
    
    # 5. 启动引擎线程
    self._start_background_thread()

初始化过程中的配置详解：

ModelConfig：包含模型相关的配置，如模型路径、类型、量化设置等

class ModelConfig:
    def __init__(
        self,
        model: str,
        tokenizer: Optional[str] = None,
        tokenizer_mode: str = "auto",
        trust_remote_code: bool = False,
        dtype: str = "auto",
        quantization: Optional[str] = None,
        revision: Optional[str] = None,
        code_revision: Optional[str] = None,
        tokenizer_revision: Optional[str] = None,
        # ... 更多参数
    ) -> None:
        # 初始化模型配置

DeviceConfig：包含设备相关配置，如GPU数量、每个GPU的最大内存等

class DeviceConfig:
    def __init__(
        self,
        device_type: Optional[str] = None,
        max_memory_per_gpu: Optional[Dict[int, str]] = None,
        # ... 更多参数
    ) -> None:
        # 初始化设备配置

ParallelConfig：包含并行计算相关配置，如张量并行度、流水线并行度等

class ParallelConfig:
    def __init__(
        self,
        tensor_parallel_size: int = 1,
        pipeline_parallel_size: int = 1,
        # ... 更多参数
    ) -> None:
        # 初始化并行配置

3.3.2 Tokenizer 初始化

Tokenizer 负责将输入文本转换为模型可以处理的 token ID 序列：

def _init_tokenizer(self) -> PreTrainedTokenizer:
    """初始化 tokenizer"""
    # 加载 tokenizer
    tokenizer = get_tokenizer(
        model_name=self.model_config.model,
        tokenizer_name=self.model_config.tokenizer,
        tokenizer_mode=self.model_config.tokenizer_mode,
        trust_remote_code=self.model_config.trust_remote_code,
        revision=self.model_config.tokenizer_revision,
    )
    
    # 配置特殊 token
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        
    # 检查并处理特殊 token ID
    if (self.model_config.chat_template is not None and
            self.model_config.chat_template != "none"):
        tokenizer.chat_template = self.model_config.chat_template
        
    return tokenizer

3.3.3 Scheduler 初始化

Scheduler 负责调度请求，决定哪些请求应该被处理以及何时处理：

def _init_schedulers(self) -> List[Scheduler]:
    """初始化调度器"""
    # 创建多个调度器实例，每个对应一个模型副本
    schedulers = []
    for rank in range(self.num_nodes):
        # 创建调度器实例
        scheduler = Scheduler(
            # 配置调度器
            scheduler_config=self.scheduler_config,
            max_num_batched_tokens=self.worker_pool.max_num_batched_tokens,
            max_num_seqs=self.worker_pool.max_num_seqs,
            device=f"cuda:{
     
     self.parallel_config.local_rank}" if torch.cuda.is_available() else "cpu",
        )
        schedulers.append(scheduler)
    return schedulers

3.3.4 Worker 池初始化

Worker 池负责管理多个 Worker 实例，处理并发请求：

class WorkerPool:
    def __init__(
        self,
        model_config: ModelConfig,
        parallel_config: ParallelConfig,
        scheduler_config: SchedulerConfig,
        device_config: DeviceConfig,
        local_rank: int,
    ) -> None:
        # 1. 设置基本配置
        self.model_config = model_config
        self.parallel_config = parallel_config
        self.scheduler_config = scheduler_config
        self.device_config = device_config
        self.local_rank = local_rank
        
        # 2. 初始化其他参数
        self.max_num_seqs = scheduler_config.max_num_seqs
        self.max_num_batched_tokens = scheduler_config.max_num_batched_tokens
        
        # 3. 初始化 Worker
        self.workers: List[Worker] = []
        self.model_replicas = []
        
        # 4. 加载模型并创建 Worker
        self._initialize_workers()

Worker 初始化时会加载模型并准备执行环境：

def _initialize_workers(self) -> None:
    """初始化 Worker"""
    # 创建模型副本
    self.model_replicas = self._initialize_model_replicas()
    
    # 为每个模型副本创建一个 Worker
    for i, model_replica in enumerate(self.model_replicas):
        worker = Worker(
            model_config=self.model_config,
            model_replica=model_replica,
            # ... 其他参数
        )
        self.workers.append(worker)

3.4 推理阶段详细流程

3.4.1 请求接收与处理

用户通过 API 发送请求，LLMEngine 接收并处理请求：

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,
) -> Optional[int]:
    """添加一个新请求"""
    # 1. 记录请求到达时间
    if arrival_time is None:
        arrival_time = time.time()
    
    # 2. 处理输入
    if prompt is None and prompt_token_ids is None:
        raise ValueError("Either prompt or prompt_token_ids must be provided")
    
    # 3. 如果需要，使用 tokenizer 处理输入文本
    if prompt_token_ids is None:
        assert prompt is not None
        prompt_token_ids = self.tokenizer.encode(prompt)
    
    # 4. 创建请求并添加到队列
    request = Request(
        request_id=request_id,
        prompt=prompt,
        prompt_token_ids=prompt_token_ids,
        sampling_params=sampling_params,
        arrival_time=arrival_time,
    )
    
    # 5. 将请求添加到请求跟踪器
    self.request_tracker.add_request(request)
    
    # 6. 将请求添加到调度器
    self.schedulers[0].add_request(request)
    
    return len(prompt_token_ids)

3.4.2 请求调度

Scheduler 负责决定何时处理哪些请求：

def schedule(self) -> List[SequenceGroupMetadata]:
    """调度请求执行"""
    # 1. 获取系统状态
    now = time.time()
    
    # 2. 从等待队列和运行中队列选择可处理的序列组
    seq_groups = self._get_schedulable_groups()
    
    # 3. 根据调度策略对序列组排序
    seq_groups = self._sort_by_strategy(seq_groups)
    
    # 4. 根据资源约束选择可处理的序列组
    scheduled_groups = self._select_best_fitting_groups(seq_groups)
    
    # 5. 生成调度元数据
    metadata = self._create_schedule_metadata(scheduled_groups)
    
    return metadata

3.4.3 模型执行

ModelRunner 负责执行模型的前向传播：

def execute_model(
    self,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    kv_caches: List[torch.Tensor],
    block_tables: torch.Tensor,
) -> Tuple[torch.Tensor, List[torch.Tensor]]:
    """执行模型前向传播"""
    # 1. 准备输入
    batch_size = input_ids.shape[0]
    
    # 2. 如果已启用 CUDA 图捕获并且批大小匹配，使用 CUDA 图执行
    if (self.model_captured and batch_size in self.graph_runners and
            not self.disable_graphs):
        # 使用预先捕获的 CUDA 图执行模型
        return self._execute_with_graph(batch_size, input_ids, positions,
                                       kv_caches, block_tables)
    else:
        # 使用标准方式执行模型
        return self._forward_helper(input_ids, positions, kv_caches,
                                    block_tables)

3.4.4 KV 缓存管理

BlockManager 负责管理 KV 缓存的物理块：

def allocate(self, seq_id: int, num_tokens: int) -> List[int]:
    """为序列分配 KV 缓存块"""
    # 1. 计算需要的块数
    num_blocks = (num_tokens + self.block_size - 1) // self.block_size
    
    # 2. 分配物理块
    block_indices = []
    for _ in range(num_blocks):
        block_idx = self.block_allocator.allocate()
        block_indices.append(block_idx)
    
    # 3. 更新映射表
    self.block_tables[seq_id] = block_indices
    
    return block_indices

3.4.5 结果处理与返回

LLMEngine 处理模型输出并返回生成结果：

def _process_model_outputs(
    self,
    seq_group: SequenceGroup,
    outputs: ModelOutput,
    metadata: SequenceGroupMetadata,
) -> None:
    """处理模型输出"""
    # 1. 处理 logits 和采样结果
    logprobs = outputs.logprobs
    next_tokens = outputs.next_tokens
    
    # 2. 为每个序列更新状态
    for seq_id, seq in seq_group.seqs.items():
        # 获取对应的 token 和概率
        next_token = next_tokens[seq_id]
        logprob = logprobs[seq_id] if logprobs is not None else None
        
        # 更新序列状态
        seq.append_token_id(next_token, logprob)
        
        # 检查是否达到终止条件
        if self._is_finished(seq, next_token):
            seq.status = SequenceStatus.FINISHED

3.5 完整调用链路示例

以下是从用户请求到生成文本的完整调用链路示例：

用户 API 调用：

from vllm import LLM, SamplingParams

# 设置采样参数
sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=100
)

# 初始化 LLM
llm = LLM(model="gpt-3.5-turbo")

# 发送请求
outputs = llm.generate("Tell me a joke", sampling_params=sampling_params)

# 获取结果
print(outputs[0].outputs[0].text)

LLM 初始化：

# LLM 类内部初始化 LLMEngine
self.engine = LLMEngine(
    model=model,
    tokenizer=tokenizer,
    # 其他参数...
)

请求处理：

# LLM.generate 方法内部调用 LLMEngine.add_request
request_id = self.engine.add_request(
    prompt=prompt,
    sampling_params=sampling_params
)

请求调度与执行：

# LLMEngine 内部的调度循环
while not engine_stopped:
    # 1. 调度请求
    scheduled_groups = self.scheduler.schedule()
    
    # 2. 将调度的请求发送给 Worker
    for worker, metadata in zip(self.workers, scheduled_groups):
        worker.execute_model(metadata)
    
    # 3. 处理模型输出
    for output in outputs:
        self._process_model_outputs(output)
    
    # 4. 更新请求状态
    self.request_tracker.update_status()

结果返回：

# LLMEngine 将结果返回给 LLM
finished_requests = self.request_tracker.get_finished()

# LLM 将结果返回给用户
return [r.get_result() for r in finished_requests]

3.6 关键调用路径总结

vLLM 的核心调用流程可以总结为以下几个关键路径：

初始化路径：

LLM.__init__ → LLMEngine.__init__ → 
[ModelConfig, DeviceConfig, ParallelConfig 初始化] → 
Tokenizer 初始化 → Worker 池初始化 → 
[模型加载, Scheduler 初始化] → 引擎线程启动

请求处理路径：

LLM.generate → LLMEngine.add_request → 
Tokenizer.encode → RequestTracker.add_request → 
Scheduler.add_request → Scheduler.schedule →
Worker.execute_model → ModelRunner.forward → 
[BlockManager 管理 KV 缓存] →
LLMEngine._process_model_outputs → 
LLM 返回结果

并行执行路径：

多个 Worker 并行执行 →
每个 Worker 内部使用 CUDA 图加速 →
使用张量并行和流水线并行提高吞吐量

通过这种模块化的设计，vLLM 能够高效地处理大量并发请求，同时优化内存使用和计算资源分配，为大型语言模型提供高性能的推理服务。

4. vLLM 关键工作机制

vLLM 的高性能源于几项关键技术创新。这些机制共同解决了大语言模型推理中的主要瓶颈问题：内存效率、计算速度和资源利用率。以下是这些核心机制的详细分析，即使没有深厚的技术背景，也能理解它们解决的问题和带来的好处。

5.1 PagedAttention 机制

PagedAttention 是 vLLM 的核心创新，通过分页管理 KV 缓存提高内存效率。这一机制受操作系统中的虚拟内存管理启发，解决了大语言模型推理中的关键瓶颈问题。

KV 缓存的挑战

在标准的 Transformer 解码器中，每生成一个新 token，都需要计算注意力，这涉及与之前所有 token 的 key 和 value 的交互。随着序列长度增加，这些 key 和 value 缓存（即 KV 缓存）占用的内存也线性增长：

# 传统 KV 缓存的内存占用
memory_per_seq = seq_len * num_layers * 2 * hidden_dim * dtype_size

对于长序列和批处理请求，内存消耗非常显著。例如，一个批次中有多个不同长度的序列，传统方法需要为每个序列分配足够大的连续内存块，导致内存碎片和浪费。

PagedAttention 原理

PagedAttention 的核心思想是将 KV 缓存组织成固定大小的物理块，并使用块表映射逻辑位置到物理块：

class PagedAttention:
    def __init__(self, block_size: int):
        self.block_size = block_size
        self.block_tables = {
   
   }  # 序列到物理块的映射
        self.physical_blocks = []  # 实际存储块
        
    def allocate_block(self, seq_id: int):
        """为序列分配新块"""
        if seq_id not in self.block_tables:
            self.block_tables[seq_id] = []
            
        # 分配新物理块
        block_id = self._find_free_block()
        self.block_tables[seq_id].append(block_id)
        
    def get_kv_cache(self, seq_id: int, positions: torch.Tensor):
        """获取缓存内容"""
        if seq_id not in self.block_tables:
            return None
            
        # 计算物理块索引和块内位置
        block_indices, block_offsets = self._compute_indices(positions)
        physical_blocks = [self.block_tables[seq_id][idx] for idx in block_indices]
        
        # 获取缓存内容
        return self._gather_blocks(physical_blocks, block_offsets)

工作流程：

内存分块：KV 缓存被划分为固定大小的块（如每块 16 个 token）
动态分配：每个序列动态分配所需的块，而不是预先分配
虚拟寻址：使用块表将逻辑位置映射到物理块
资源复用：当序列完成时，其块可以被释放并重用

技术细节

实际实现中，PagedAttention 包含以下关键组件：

BlockAllocator：管理物理块的分配和释放

class BlockAllocator:
    def __init__(self, num_blocks: int):
        self.num_blocks = num_blocks
        self.free_blocks = set(range(num_blocks))
        self.used_blocks = set()
        
    def allocate(self) -> int:
        """分配一个空闲块"""
        if not self.free_blocks:
            raise RuntimeError("No free blocks available")
        
        block_id = next(iter(self.free_blocks))
        self.free_blocks.remove(block_id)
        self.used_blocks.add(block_id)
        return block_id
        
    def free(self, block_id: int) -> None:
        """释放一个块"""
        self.used_blocks.remove(block_id)
        self.free_blocks.add(block_id)

BlockTable：维护序列到物理块的映射

class BlockTable:
    def __init__(self, block_size: int):
        self.block_size = block_size
        self.tables = {
   
   }  # seq_id -> List[block_id]
        
    def add_block(self, seq_id: int, block_id: int) -> None:
        """为序列添加一个块"""
        if seq_id not in self.tables:
            self.tables[seq_id] = []
        self.tables[seq_id].append(block_id)
        
    def get_physical_blocks(self, seq_id: int, positions: torch.Tensor) -> Tuple[List[int], torch.Tensor]:
        """获取物理块 ID 和偏移量"""
        block_ids = self.tables[seq_id]
        block_indices = positions // self.block_size
        block_offsets = positions % self.block_size
        physical_blocks = [block_ids[idx.item()] for idx in block_indices]
        return physical_blocks, block_offsets

注意力计算核心：高效实现分块注意力计算

# CUDA 核心伪代码
@cuda.kernel
def paged_attention_kernel(
    q: Tensor,             # [batch_size, num_heads, head_dim]
    k_cache: Tensor,       # [num_blocks, block_size, num_heads, head_dim]
    v_cache: Tensor,       # [num_blocks, block_size, num_heads, head_dim]
    block_tables: Tensor,  # [batch_size, max_blocks]
    output: Tensor         # [batch_size, num_heads, head_dim]
):
    # 获取线程索引
    batch_idx = cuda.blockIdx.x
    head_idx = cuda.blockIdx.y
    
    # 获取该序列的块表
    seq_blocks = block_tables[batch_idx]
    
    # 本地计算缓冲区
    local_k = shared_memory[...]
    local_v = shared_memory[...]
    
    # 加载查询向量
    query = q[batch_idx, head_idx]
    
    # 对每个块执行注意力计算
    for i in range(len(seq_blocks)):
        block_id = seq_blocks[i]
        
        # 加载 KV 缓存到共享内存
        local_k.copy_from(k_cache[block_id])
        local_v.copy_from(v_cache[block_id])
        
        # 计算注意力分数
        scores = compute_attention(query, local_k)
        
        # 应用注意力权重
        output[batch_idx, head_idx] += apply_attention(scores, local_v)

性能优势

PagedAttention 带来了多项显著优势：

内存效率：通过分块和动态分配，显著减少内存碎片
- 传统方法：为每个序列分配最大长度的连续内存
- PagedAttention：按需分配块，高效共享物理内存
批处理吞吐量：支持更多并发请求
- 实验表明，相同硬件上，vLLM 可以处理比其他框架多 2-4 倍的并发请求
长文本支持：优雅处理长序列生成
- 只需添加新块，无需重新分配或复制整个缓存
内存利用率：内存使用量与实际 token 数成正比，而非预分配
- 典型场景中，与传统方法相比，内存使用可减少 50-70%

实际应用案例

在 vLLM 中，PagedAttention 的应用可以在多方面观察到：

动态批处理：高效支持连续批处理（Continuous Batching）

# 在每次前向传播中添加新序列
def step(self):
    # 调度新序列和正在进行的序列
    scheduled_seq_groups = self.scheduler.schedule()
    if not scheduled_seq_groups:
        return
    
    # 准备批处理输入
    batch = self._prepare_batch(scheduled_seq_groups)
    
    # 为新序列分配 KV 缓存块
    for seq_id in batch.new_seq_ids:
        self.block_manager.allocate_blocks(seq_id, batch.prompt_lens[seq_id])
        
    # 执行模型前向传播
    outputs = self.model_runner.forward(batch.input_ids, batch.block_tables, ...)

高效内存管理：动态分配和回收物理块

def free_finished_sequences(self):
    """释放已完成序列的内存"""
    for seq_id in self.finished_sequences:
        # 获取分配的块
        blocks = self.block_tables.get_blocks(seq_id)
        
        # 释放块
        for block_id in blocks:
            self.block_allocator.free(block_id)
            
        # 移除块表条目
        self.block_tables.remove(seq_id)

缓存共享：支持 fork 和复制操作

def fork_sequence(self, src_seq_id: int, dst_seq_id: int, fork_pos: int

本文章已经生成可运行项目

最低0.47元/天解锁文章