🦞 一只用 AI Agent 搭副业产线的程序员


你跑过一个 70B 模型吗?Q4 量化后大概 40GB,一张 A100(80GB)放得下。但生产环境的问题不是"放不放得下",而是一个请求只用了 2K 上下文,为什么显存就不够处理第二个请求了

答案是 KV Cache。传统方法预分配了太多永远用不到的显存空间。

vLLM 的 PagedAttention 就是来解决这个问题的。这篇文章我们看它怎么把操作系统的虚拟内存管理思想搬到了 GPU 显存管理上。


项目简介

vLLM(GitHub 40k+ Stars)是 UC Berkeley 开源的 LLM 推理引擎,核心贡献是 PagedAttention——一种把 KV Cache 按"页"管理的算法。它把显存利用率从传统框架的 30-40% 提升到 90% 以上,吞吐量提升 2-4 倍。现在被 LMSYS(Chatbot Arena)和多家公司用于生产环境。


架构全景

┌──────────────────────────────────────────────────────────────┐
│                       API 服务层                              │
│       OpenAI-compatible: /v1/completions, /v1/chat,         │
│                          /v1/embeddings, /v1/models          │
├──────────────────────────────────────────────────────────────┤
│                     调度器(Scheduler)                        │
│   Continuous Batching — 不再等整批完成,来一个处理一个          │
│   ┌─────────┐  ┌──────────┐  ┌───────────────┐              │
│   │ 请求队列  │→│ Prefill  │→│    Decode 循环   │              │
│   └─────────┘  └──────────┘  └───────────────┘              │
├──────────────────────────────────────────────────────────────┤
│              块管理器(Block Manager)—— PagedAttention 核心   │
│   ┌──────────┐  ┌───────────┐  ┌─────────────┐             │
│   │  物理块池  │  │  块映射表  │  │  Copy-on-Write │           │
│   │ (Physical) │  │ (BlockTable)│  │  (beam search) │          │
│   └──────────┘  └───────────┘  └─────────────┘             │
├──────────────────────────────────────────────────────────────┤
│                    CUDA Kernel 层                             │
│   PagedAttention · FlashAttention · FP8/INT8 量化 · TP/PP    │
└──────────────────────────────────────────────────────────────┘

先理解问题:传统 KV Cache 为什么浪费显存

LLM 推理时,每生成一个 token,都要拿当前的 query 去和之前所有 token 的 key/value 做 attention。为了避免重复计算,程序把每一层的 K 和 V 张量存下来——这就是 KV Cache。

传统框架(FasterTransformer、TGI)的处理方式:

请求 1,上下文 2000 tokens → 预分配 (max_context=4096) × K × V ≈ 2GB
请求 2,上下文 500 tokens → 预分配 (max_context=4096) × K × V ≈ 2GB
请求 3,上下文 8000 tokens → 预分配 (max_context=8192) × K × V ≈ 4GB
─────────────────────────────────────────────────────────────
总占用:8GB(但实际只用到了 (2000+500+8000)/(4096+4096+8192) ≈ 64%)

两个问题:

  1. 预分配:必须按"最大可能长度"分配,绝大多数请求用不完。
  2. 碎片化:请求 1 结束后释放 2GB,但紧接着来一个需要 3GB 的请求——那 2GB 的碎片用不上,要 defrag 或者 OOM。

核心问题:KV Cache 是连续分配的。连续分配 = 外部碎片 = 浪费。


关键设计一:分页管理——显存版的虚拟内存

vLLM 的答案是:不按"请求"分配,按"块"分配。每个块固定大小(比如 16 个 token),物理块池化,请求通过"块表"引用物理块。

# vllm/core/block_manager.py —— 块管理器的核心逻辑(概念性重建)
from typing import List, Optional, Dict

class BlockTable:
    """每个请求的虚拟块 → 物理块的映射表"""
    def __init__(self, block_size: int = 16):
        self.block_size = block_size              # 每个物理块 = 16 个 token
        self.blocks: List[Optional[int]] = []     # 虚拟块号 → 物理块号

class BlockAllocator:
    """全局物理块池"""
    def __init__(self, num_blocks: int, block_size: int):
        self.free_blocks: List[int] = list(range(num_blocks))  # 空闲块列表
        self.block_size = block_size

    def allocate(self) -> int:
        """分配一个物理块,返回块号"""
        if not self.free_blocks:
            raise OutOfMemoryError("No free blocks")
        return self.free_blocks.pop()

    def free(self, block_id: int):
        """释放物理块"""
        self.free_blocks.append(block_id)

class BlockManager:
    """全局块管理器——所有请求共享物理块池"""
    def __init__(self, num_gpu_blocks: int, block_size: int = 16):
        self.allocator = BlockAllocator(num_gpu_blocks, block_size)
        self.block_tables: Dict[int, BlockTable] = {}  # 请求 ID → 块表

    def append_slot(self, seq_id: int) -> Optional[int]:
        """为一个请求追加一个 slot(需要时分配新块)"""
        block_table = self.block_tables[seq_id]
        # 计算需要几个块
        num_needed = (len(block_table.blocks) * self.allocator.block_size) + 1
        # 如果最后一个块已满,分配新块
        if num_needed > len(block_table.blocks) * self.allocator.block_size:
            new_block = self.allocator.allocate()
            block_table.blocks.append(new_block)
        # 返回最后一个物理块的地址
        return block_table.blocks[-1]

这个设计的效果:

改前(连续分配):
请求A: ┌──────────4096 tokens──────────┬ 碎片 ┐
请求B: ┌────2048 tokens────┬ 碎片 ┐

改后(分页分配,块大小=16):
物理块池: [A1][B1][A2][空][B2][A3][空][空][A4][B3]...
         └── A 的块表: [0, 2, 5, 8] ──┘
                └── B 的块表: [1, 4, 9] ──┘

没有外部碎片了——因为所有分配都是固定大小的块。内部碎片最多 15 个 token(最后一个块没装满),在上下文的尺度下可以忽略不计。

设计洞察:这就是操作系统的分页思想,直接搬到 GPU 显存管理。页表 + 物理页框池 + 按需分配。香不香?香。新不新?不新。但能把 60 年前的 OS 思想用到 LLM 推理里并做到生产可用——这就是工程的魅力。


关键设计二:Copy-on-Write——并行生成的零拷贝优化

一个常见场景:用户要求"生成 3 个候选回复"。怎么做?

朴素方案:KV Cache 复制 3 份。一个 4K 上下文的请求 = 2GB KV Cache。3 份 = 6GB。

PagedAttention 方案:共享前缀部分的物理块,只在分叉点复制块表指针。

# vllm/core/block_manager.py —— Copy-on-Write fork
class BlockManager:
    def fork(self, parent_seq_id: int, child_seq_id: int):
        """从父请求 fork 一个子请求(beam search / parallel sampling)"""
        parent_table = self.block_tables[parent_seq_id]
        # 子请求共享父请求的块表(shallow copy)
        child_table = BlockTable(block_size=parent_table.block_size)
        child_table.blocks = list(parent_table.blocks)  # 引用相同的物理块

        self.block_tables[child_seq_id] = child_table

    def append_slot(self, seq_id: int) -> Optional[int]:
        """追加 slot——如果物理块被共享,先 Copy-on-Write"""
        block_table = self.block_tables[seq_id]
        last_block = block_table.blocks[-1] if block_table.blocks else None

        # 检查最后一个块是否被多个请求共享
        if last_block is not None and self._ref_count(last_block) > 1:
            # COW: 分配新物理块,复制内容
            new_block = self.allocator.allocate()
            self._copy_block(last_block, new_block)
            self.allocator.free(last_block)  # 减少旧块的引用计数
            block_table.blocks[-1] = new_block

        # 剩下的逻辑跟普通 append 一样
        ...

这个优化让 parallel sampling(生成 n 个候选回复)的显存开销从 O(n) 降到 O(1),只额外花在分叉后产生差异的 token 上。

设计洞察:Copy-on-Write 的通用性极高——fork 进程用它、Redis 的 BGSAVE 用它、vLLM 的 parallel sampling 也用它。理解一个模式,能用一辈子。


关键设计三:Continuous Batching——请求级别的流水线

传统批处理(Static Batching):等到一批请求全部完成,再处理下一批。

请求A(200 tokens)→ ████████████████████
请求B(50 tokens) →  █████ → 等 A 完成 → 空闲
请求C(10 tokens) →  ██ → 等 A 和 B 完成 → 空闲

Continuous Batching:一个请求完成立即踢出,把空出来的计算资源给等待队列的下一个。

# vllm/core/scheduler.py —— 调度器的核心逻辑(概念性重建)
class Scheduler:
    def schedule(self) -> SchedulerOutput:
        running: List[SequenceGroup] = []
        preempted: List[SequenceGroup] = []

        # Step 1: 从等待队列拉请求,直到显存不够
        while self.waiting and self.block_manager.can_allocate():
            seq_group = self.waiting.pop(0)
            self.block_manager.allocate(seq_group)
            running.append(seq_group)

        # Step 2: 为每个运行中的请求生成一个 token
        for seq_group in running:
            seq_group.generate_one_token()

        # Step 3: 把完成的请求踢出,释放块
        for seq_group in running:
            if seq_group.is_finished():
                self.block_manager.free(seq_group)
                running.remove(seq_group)

        # Step 4: 剩余请求继续下一轮调度
        return SchedulerOutput(
            scheduled=running,
            preempted=preempted,
            num_waiting=len(self.waiting),
        )

调度的核心在 Step 1:“能分配就分配”。不等到"最佳批次大小",而是只要显存有空就拉新请求。

这个设计的关键收益:短请求不用等长请求。50 个 token 的请求,生成完立刻释放 KV Cache 块,给下一个请求腾空间。在混合长短请求的场景下,吞吐量提升最明显。


核心代码拆解:PagedAttention 的 CUDA Kernel 是怎么读取 KV Cache 的

把 KV Cache 分页之后,attention 计算就不能用连续的矩阵乘法了——K 和 V 分散在不同物理块里。vLLM 为此写了一个定制的 CUDA kernel:

// vllm/csrc/attention/paged_attention.cu —— 简化逻辑
__global__ void paged_attention_kernel(
    float* output,           // [num_tokens, num_heads, head_size]
    const float* query,      // [num_tokens, num_heads, head_size]
    const float* key_cache,  // [num_blocks, num_heads, block_size, head_size]
    const float* value_cache,// [num_blocks, num_heads, block_size, head_size]
    const int* block_table,  // [num_requests, max_num_blocks]
    const int* context_lens, // [num_requests]
    int   num_heads,
    float scale,
    int   block_size
) {
    int tid = threadIdx.x;
    int seq_idx = blockIdx.x;  // 每个请求一个 block
    int head_idx = blockIdx.y; // 每个 head 一个…嗯,另一个 block

    int num_blocks = (context_lens[seq_idx] + block_size - 1) / block_size;

    for (int block_idx = 0; block_idx < num_blocks; block_idx++) {
        // 关键:通过块表把"虚拟块号"转成"物理块号"
        int physical_block = block_table[seq_idx * max_num_blocks + block_idx];

        // 用物理块号去读 K 和 V
        // key_cache[physical_block * block_stride + head_idx * block_size * head_size + ...]
        int block_offset = physical_block * block_stride
                         + head_idx * block_size * head_size;

        // 计算这个 token 和当前块里所有 token 的 attention score
        for (int t = 0; t < block_size; t++) {
            float score = 0;
            for (int d = 0; d < head_size; d++) {
                score += query[q_offset + d] * key_cache[block_offset + t * head_size + d];
            }
            scores[block_idx * block_size + t] = score * scale;
        }
    }

    // softmax + weighted sum(跟标准 attention 一样)
    ...
}

kernel 的核心只有一行physical_block = block_table[seq_idx * max_blocks + block_idx]。这一行就是 PagedAttention 的全部魔法——其余部分都是在做普通的 attention 运算。代价只是一次额外的全局内存读取(读块表),在已经在做 O(n^2) 的 attention 运算面前可以忽略。


你可以抄的作业

1. 分页思路不只属于操作系统

任何"大块连续分配会碎片化"的场景都能用分页。你要做个内存池管理游戏对象?分页。管理网络数据包的缓冲?分页。只要分配单元不固定且碎片化严重,固定大小的块 + 映射表就是标准答案。

2. Copy-on-Write 的通用模式

fork + 共享读取 + 写时复制。vLLM 的 block fork 和 Linux 的 fork() 系统调用本质上是一回事。理解 COW,你就多了一个优化"同源并行操作"的武器。

3. Continuous Batching 就是"能进就进"

Pipeline 优化的本质不是"凑满批次",而是"不要等"。任何批处理系统——ETL 管道、消息队列消费者、API Gateway——都可以用这个思路:有空闲资源就拉,做完立刻放,不等。

4. 好的优化 = 好的数据结构 + 老的计算机思想

PagedAttention 没有发明新数学,它只是选了正确的数据结构(页表)和正确的资源管理策略(按需分配 + COW)。“创新"很多时候就是"把别处已经验证过的思想搬到你的领域”。


最后

vLLM 的故事跟很多"学术项目变成工业标准"的故事一样:不是算法有多新,而是工程优化做得够深。PagedAttention 把显存利用率从 30% 拉到 90%,靠的不是一篇论文、一个公式,而是一个完整的系统设计——块管理、调度器、CUDA kernel 三者协同工作。

理解 vLLM 不只是学一个推理框架,更是学"如何把一个孤立的算法优化做成系统级的解决方案"。

下一讲拆 Dify。一个低代码 AI 应用平台,是怎么设计插件系统让 100+ 个模型和工具能无缝接入的?


本文拆解的 vLLM 版本:v0.6.x。源码地址:github.com/vllm-project/vllm


🦞 一只用 AI Agent 搭副业产线的程序员

全平台同名:虾哥不加班 | 源码:GitHub - lobster-bujiaban
需要定制 AI 工具?来聊聊 → lob_ai

Logo

免费领 150 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐