在大语言模型(LLM)的时代,推理效率成为了制约应用落地的关键瓶颈。想象一下,当你向 ChatGPT 提出问题时,它需要在几秒内处理数千个token,同时服务数万名用户。这背后的技术挑战是什么?答案就是高效的推理引擎。

vLLM 作为当前最高效的 LLM 推理引擎之一,通过创新的 PagedAttention、KV Cache 管理、连续批处理等技术,将推理吞吐量提升了 24倍。然而,vLLM 的生产代码复杂庞大,包含大量工程优化和边界情况处理,对于初学者来说犹如迷宫。

cleanvllm 项目应运而生——它承继了 CleanRL 的教育理念,用单文件(1500行Python代码)实现了 vLLM 的核心概念,去除了工程复杂性,让我们能够专注于理解架构本质。正如 Linus Torvalds 所说:”Talk is cheap. Show me the code.”

本文将带你从零开始,逐步构建一个完整的 vLLM 推理引擎,揭示其背后的设计哲学和技术细节。读完本文,你将能够:

  • 理解 vLLM 的核心架构和关键创新
  • 掌握 PagedAttention 的工作原理
  • 学会如何实现高效的 KV Cache 管理
  • 了解连续批处理的调度策略
  • 运行自己的 vLLM 实现

vLLM 的核心挑战与设计思路

传统 LLM 推理的痛点

让我们以一个具体的例子来说明:假设你有一台 A100 GPU(80GB显存),要同时处理 100 个用户请求,每个请求平均生成 100 个token。

传统的推理方式会遇到以下问题: 1. 内存碎片化:每个序列的 KV Cache 需要连续内存,但序列长度不同导致严重的内存浪费。实际可用内存只有 60%! 2. 批处理困难:序列A需要10个token,序列B需要100个token,如何高效批处理?传统方法只能等最短序列完成。 3. 动态调度复杂:你无法预知用户会问多长的问题,资源调度变得极其困难。 这些问题导致 GPU 利用率通常只有 30-40%,大量算力被浪费。

vLLM 的解决方案

vLLM 通过三个核心创新解决了这些问题:

  • PagedAttention:将 KV Cache 分页管理,类似操作系统的虚拟内存
  • 连续批处理:动态添加/移除序列,最大化 GPU 利用率
  • 前缀缓存:复用相同前缀的 KV Cache,减少重复计算 让我们在 cleanvllm 中看看这些概念是如何实现的。

数据结构设计 - vLLM 的基石

序列管理:Sequence 类

在 cleanvllm 中,每个用户请求被抽象为一个 Sequence 对象:

class Sequence:
    """
    序列类
    
    表示一个文本生成请求的序列,包含prompt tokens和生成的completion tokens。
    管理序列的状态、token列表、块表(block table)等信息。
    """
    block_size = 256  # KV缓存的块大小
    counter = count()  # 序列ID计数器

    def __init__(self, token_ids: list[int], sampling_params = SamplingParams()):
        """
        初始化序列
        
        Args:
            token_ids: prompt的token ID列表
            sampling_params: 采样参数
        """
        self.seq_id = next(Sequence.counter)  # 序列的唯一ID
        self.status = SequenceStatus.WAITING  # 初始状态为等待
        self.token_ids = copy(token_ids)  # token ID列表的副本
        self.last_token = token_ids[-1]  # 最后一个token
        self.num_tokens = len(self.token_ids)  # token总数
        self.num_prompt_tokens = len(token_ids)  # prompt token数量
        self.num_cached_tokens = 0  # 已缓存的token数量(用于前缀缓存)
        self.block_table = []  # 关键:**块表**,存储KV缓存块的索引, 映射逻辑块到物理块
        self.temperature = sampling_params.temperature  # 采样温度
        self.max_tokens = sampling_params.max_tokens  # 最大生成token数
        self.ignore_eos = sampling_params.ignore_eos  # 是否忽略EOS token

    def __len__(self):
        """返回序列的token数量"""
        return self.num_tokens

    def __getitem__(self, key):
        """支持索引访问token列表"""
        return self.token_ids[key]

    #@property 的核心作用: 把“计算得出的值”封装成只读属性,让接口更自然、更安全
    # 调用方式是:seq.is_finished      # ✅ 像属性一样访问
    #     而不是: seq.is_finished()    # ❌ 不需要加括号
    @property
    def is_finished(self):
        """检查序列是否已完成"""
        return self.status == SequenceStatus.FINISHED

    @property
    def num_completion_tokens(self):
        """返回已生成的completion token数量"""
        return self.num_tokens - self.num_prompt_tokens

    @property
    def prompt_token_ids(self):
        """返回prompt部分的token ID列表"""
        return self.token_ids[:self.num_prompt_tokens]

    @property
    def completion_token_ids(self):
        """返回completion部分的token ID列表"""
        return self.token_ids[self.num_prompt_tokens:]

    @property
    def num_cached_blocks(self):
        """返回已缓存的块数量"""
        return self.num_cached_tokens // self.block_size

    @property
    def num_blocks(self):
        """返回序列需要的总块数量"""
        return (self.num_tokens + self.block_size - 1) // self.block_size

    @property
    def last_block_num_tokens(self):
        """返回最后一个块中的token数量"""
        return self.num_tokens - (self.num_blocks - 1) * self.block_size

    def block(self, i):
        """
        获取第i个块的token列表
        
        Args:
            i: 块索引
            
        Returns:
            该块的token ID列表
        """
        assert 0 <= i < self.num_blocks
        return self.token_ids[i*self.block_size: (i+1)*self.block_size]

    def append_token(self, token_id: int):
        """
        向序列追加一个新生成的token
        
        Args:
            token_id: 新生成的token ID
        """
        self.token_ids.append(token_id)
        self.last_token = token_id
        self.num_tokens += 1

    def __getstate__(self):
        """
        序列化函数,用于pickle
        
        为了节省内存,如果序列还没有生成completion tokens,则保存完整的token列表;
        否则只保存最后一个token(因为prompt tokens可以从其他地方恢复)
        """
        return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table,
                self.token_ids if self.num_completion_tokens == 0 else self.last_token)

    def __setstate__(self, state):
        """
        反序列化函数,用于pickle
        
        从序列化状态恢复序列对象
        """
        self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table = state[:-1]
        if self.num_completion_tokens == 0:
            self.token_ids = state[-1]
        else:
            self.last_token = state[-1]

关键洞察:block_table 是 PagedAttention 的核心数据结构,类似操作系统的页表。它将序列的逻辑内存映射到物理内存块,实现了内存的灵活管理。

执行上下文:Context 类

"""
上下文管理模块

实现了全局上下文管理,用于在模型推理过程中传递序列信息、KV缓存位置等。
使用全局变量存储上下文,避免在函数间传递大量参数。
"""
from dataclasses import dataclass
import torch

#@dataclass: 少写样板代码(boilerplate),让“装数据的类”更干净
@dataclass
class Context:
    """
    上下文数据类
    
    存储当前推理步骤的上下文信息,包括:
    - 是否为prefill阶段
    - 序列长度信息
    - KV缓存的slot映射
    - 块表信息
    """
    is_prefill: bool = False  # 是否为prefill阶段
    cu_seqlens_q: torch.Tensor | None = None  # query的累积序列长度
    cu_seqlens_k: torch.Tensor | None = None  # key的累积序列长度
    max_seqlen_q: int = 0  # query的最大序列长度
    max_seqlen_k: int = 0  # key的最大序列长度
    slot_mapping: torch.Tensor | None = None  # KV缓存的slot映射
    context_lens: torch.Tensor | None = None  # 每个序列的上下文长度(decode阶段使用)
    block_tables: torch.Tensor | None = None  # 块表(用于前缀缓存和decode阶段)

_CONTEXT = Context()  # 全局上下文对象

def get_context():
    """
    获取全局上下文
    
    Returns:
        当前上下文对象
    """
    return _CONTEXT

def set_context(is_prefill, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=0, max_seqlen_k=0, slot_mapping=None, context_lens=None, block_tables=None):
    """
    设置全局上下文
    
    Args:
        is_prefill: 是否为prefill阶段
        cu_seqlens_q: query的累积序列长度
        cu_seqlens_k: key的累积序列长度
        max_seqlen_q: query的最大序列长度
        max_seqlen_k: key的最大序列长度
        slot_mapping: KV缓存的slot映射
        context_lens: 每个序列的上下文长度
        block_tables: 块表
    """
    global _CONTEXT
    _CONTEXT = Context(is_prefill, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, context_lens, block_tables)

def reset_context():
    """重置全局上下文"""
    global _CONTEXT
    _CONTEXT = Context()

Context 存储当前推理步骤的全局状态,区分 Prefill(首次处理提示)和 Decode(逐token生成)两个阶段。

配置管理:Config 类

"""
配置管理模块

定义了nano-vllm系统的配置类,包含模型路径、批处理参数、GPU内存使用等配置项
"""
import os
from dataclasses import dataclass
from transformers import AutoConfig


@dataclass
class Config:
    """
    系统配置类
    
    包含nano-vllm运行所需的所有配置参数,包括模型路径、批处理大小、
    GPU内存使用率、张量并行大小等。
    """
    model: str  # 模型路径
    max_num_batched_tokens: int = 16384  # 批处理的最大token数量
    max_num_seqs: int = 512  # 最大序列数量
    max_model_len: int = 4096  # 模型支持的最大序列长度
    gpu_memory_utilization: float = 0.9  # GPU内存使用率(0-1之间)
    tensor_parallel_size: int = 1  # 张量并行大小,支持1-8
    enforce_eager: bool = False  # 是否强制使用eager模式(不使用CUDA graph)
    hf_config: AutoConfig | None = None  # HuggingFace配置对象
    eos: int = -1  # 结束符token ID
    kvcache_block_size: int = 256  # KV缓存块大小,必须是256的倍数
    num_kvcache_blocks: int = -1  # KV缓存块数量,-1表示自动计算

    def __post_init__(self):
        """
        初始化后处理函数
        
        验证配置参数的有效性,并加载HuggingFace配置
        """
        # 验证模型路径是否存在
        assert os.path.isdir(self.model)
        # 验证KV缓存块大小是256的倍数
        assert self.kvcache_block_size % 256 == 0
        # 验证张量并行大小在有效范围内
        assert 1 <= self.tensor_parallel_size <= 8
        # 加载HuggingFace配置
        self.hf_config = AutoConfig.from_pretrained(self.model)
        # 确保最大模型长度不超过配置中的最大位置嵌入
        self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)
        # 确保批处理token数量至少等于最大模型长度
        assert self.max_num_batched_tokens >= self.max_model_len

这些参数控制着系统的内存使用、并发能力和性能特征。

内存管理 - PagedAttention 的实现

内存块抽象

"""
块管理器模块

实现了KV缓存的块管理功能,包括块的分配、释放、哈希缓存等。
支持前缀缓存(prefix caching)以提高性能。
"""
from collections import deque
import xxhash
import numpy as np

from nanovllm.engine.sequence import Sequence


class Block:
    """
    KV缓存块类
    
    表示一个KV缓存块,包含块ID、引用计数、哈希值和token列表。
    """

    def __init__(self, block_id):
        """
        初始化块
        
        Args:
            block_id: 块的唯一ID
        """
        self.block_id = block_id  # 块的唯一ID
        self.ref_count = 0  # 引用计数,表示有多少个序列在使用这个块
        self.hash = -1  # 块的哈希值,用于前缀缓存查找
        self.token_ids = []  # 块中存储的token ID列表

    def update(self, hash: int, token_ids: list[int]):
        """
        更新块的内容和哈希值
        
        Args:
            hash: 块的哈希值
            token_ids: token ID列表
        """
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):
        """重置块的状态"""
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []


## 块管理器:PagedAttention 的核心
class BlockManager:
    """
    块管理器类
    
    管理KV缓存块的分配和释放,支持前缀缓存以提高性能。
    通过哈希表快速查找已缓存的块,避免重复计算。
    """

    def __init__(self, num_blocks: int, block_size: int):
        """
        初始化块管理器
        
        Args:
            num_blocks: 总块数量
            block_size: 每个块的大小(token数量)
        """
        self.block_size = block_size  # 块大小
        self.blocks: list[Block] = [Block(i) for i in range(num_blocks)]  # 所有块的列表
        self.hash_to_block_id: dict[int, int] = dict()  # 哈希值到块ID的映射,用于前缀缓存
        self.free_block_ids: deque[int] = deque(range(num_blocks))  # 空闲块ID队列
        self.used_block_ids: set[int] = set()  # 已使用块ID集合

    @classmethod
    def compute_hash(cls, token_ids: list[int], prefix: int = -1):
        """
        计算token列表的哈希值
        
        Args:
            token_ids: token ID列表
            prefix: 前缀哈希值,用于链式哈希计算
            
        Returns:
            计算得到的哈希值
        """
        h = xxhash.xxh64()
        if prefix != -1:
            h.update(prefix.to_bytes(8, "little"))
        h.update(np.array(token_ids).tobytes())
        return h.intdigest()

    def _allocate_block(self, block_id: int) -> Block:
        """
        分配一个块(内部方法)
        
        Args:
            block_id: 要分配的块ID
            
        Returns:
            分配得到的块对象
        """
        block = self.blocks[block_id]
        assert block.ref_count == 0
        block.reset()
        self.free_block_ids.remove(block_id)
        self.used_block_ids.add(block_id)
        return self.blocks[block_id]

    def _deallocate_block(self, block_id: int) -> Block:
        """
        释放一个块(内部方法)
        
        Args:
            block_id: 要释放的块ID
        """
        assert self.blocks[block_id].ref_count == 0
        self.used_block_ids.remove(block_id)
        self.free_block_ids.append(block_id)

    def can_allocate(self, seq: Sequence) -> bool:
        """
        检查是否有足够的空闲块来分配序列
        
        Args:
            seq: 要检查的序列
            
        Returns:
            如果有足够块则返回True,否则返回False
        """
        return len(self.free_block_ids) >= seq.num_blocks

    def allocate(self, seq: Sequence):
        """
        为序列分配KV缓存块
        
        支持前缀缓存:如果块的token序列已经在缓存中,则复用该块,避免重复计算。
        
        Args:
            seq: 要分配块的序列
        """
        assert not seq.block_table
        h = -1  # 当前块的哈希值
        cache_miss = False  # 是否发生缓存未命中
        for i in range(seq.num_blocks):
            token_ids = seq.block(i)
            # 只有当块是完整大小(block_size)时才计算哈希值
            h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1
            # 尝试从哈希表中查找已存在的块
            block_id = self.hash_to_block_id.get(h, -1)
            # 检查是否真的匹配(哈希碰撞检查)
            if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
                cache_miss = True
            if cache_miss:
                # 缓存未命中,分配新块
                block_id = self.free_block_ids[0]
                block = self._allocate_block(block_id)
            else:
                # 缓存命中,复用已有块
                seq.num_cached_tokens += self.block_size
                if block_id in self.used_block_ids:
                    # 块已被使用,增加引用计数
                    block = self.blocks[block_id]
                    block.ref_count += 1
                else:
                    # 块未被使用,分配它
                    block = self._allocate_block(block_id)
            # 如果计算了哈希值,更新块并记录到哈希表
            if h != -1:
                block.update(h, token_ids)
                self.hash_to_block_id[h] = block_id
            seq.block_table.append(block_id)

    def deallocate(self, seq: Sequence):
        """
        释放序列占用的所有块
        
        Args:
            seq: 要释放块的序列
        """
        # 从后往前遍历块表,减少引用计数
        for block_id in reversed(seq.block_table):
            block = self.blocks[block_id]
            block.ref_count -= 1
            # 如果引用计数为0,释放该块
            if block.ref_count == 0:
                self._deallocate_block(block_id)
        seq.num_cached_tokens = 0
        seq.block_table.clear()

    def can_append(self, seq: Sequence) -> bool:
        """
        检查是否可以追加新token到序列(是否需要新块)
        
        Args:
            seq: 要检查的序列
            
        Returns:
            如果可以追加则返回True
        """
        # 如果序列长度模block_size等于1,说明需要新块
        return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

    def may_append(self, seq: Sequence):
        """
        为序列追加新token做准备(如果需要新块则分配)
        
        当序列追加新token时,可能需要:
        1. 分配新块(如果当前块已满)
        2. 更新最后一个块的哈希值(如果块刚满)
        
        Args:
            seq: 要追加token的序列
        """
        block_table = seq.block_table
        last_block = self.blocks[block_table[-1]]
        if len(seq) % self.block_size == 1:
            # 需要新块:当前块已满,需要分配新块
            assert last_block.hash != -1
            block_id = self.free_block_ids[0]
            self._allocate_block(block_id)
            block_table.append(block_id)
        elif len(seq) % self.block_size == 0:
            # 块刚满:更新最后一个块的哈希值
            assert last_block.hash == -1
            token_ids = seq.block(seq.num_blocks-1)
            # 使用前一个块的哈希值作为前缀
            prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1
            h = self.compute_hash(token_ids, prefix)
            last_block.update(h, token_ids)
            self.hash_to_block_id[h] = last_block.block_id
        else:
            # 块未满:不需要操作
            assert last_block.hash == -1

每个 Block 代表一个固定大小(通常256个token)的内存块,可以存储 KV Cache。hash 字段支持前缀缓存优化。

微信扫一扫

作者:mospan
微信关注:墨斯潘園
本文出处:http://mospany.github.io/2025/12/14/nano-vllm-code-analysis/
文章版权归本人所有,欢迎转载,但未经作者同意必须保留此段声明,且在文章页面明显位置给出原文连接,否则保留追究法律责任的权利。