vLLM Deep Dive Part 2: PagedAttention - The Core Innovation

Posted on Edited on In

第二部分：PagedAttention——核心创新

简介

PagedAttention 是 vLLM 的突破性创新，彻底改变了大语言模型的服务方式。在 PagedAttention 出现之前，高效服务 LLM 长期受到内存碎片化和浪费问题的困扰。本文将深入探讨 PagedAttention 的工作原理、重要意义，以及它在 vLLM 中的实现方式。

LLM 服务中的内存问题

理解 KV Cache

在 transformer 模型中，attention 层为每个 token 计算键（K）和值（V）向量。在生成过程中：

  1. Prefill 阶段：为所有提示词 token 计算 K、V
  2. Decode 阶段：对每个新 token，计算其 K、V，并对所有之前的 K、V 进行 attention 计算

挑战：我们必须存储所有之前的 K、V 张量（即”KV cache”），以避免重复计算。

内存大小：以 Llama-3-8B 为例：

  • 32 个 attention 层
  • 4096 隐藏维度
  • FP16 精度（2 字节）
  • 每个 token：32 层 × 4096 维度 × 2（K、V）× 2 字节 = 512 KB

对于 2048 个 token 的序列：仅 KV cache 就需要 1 GB！

碎片化问题

传统服务系统为每个请求的最大长度预先分配连续内存：

1
2
3
4
5
6
7
8
Request 1 (max 2048 tokens, actual 512):
[████░░░░░░░░░░░░] 75% wasted

Request 2 (max 2048 tokens, actual 1800):
[████████████████] 12% wasted

Request 3 (max 2048 tokens, actual 200):
[██░░░░░░░░░░░░░░] 90% wasted

问题：

  1. 过度分配：必须按最坏情况的长度进行分配
  2. 碎片化：无法将一个请求的未使用空间分配给另一个请求
  3. 低吞吐量：内存成为瓶颈，而非计算

实际影响：没有 PagedAttention，你可能只能服务 10 个并发请求；而有了 PagedAttention，同样的内存可以服务 50 个以上的请求！

PagedAttention 的解决方案

PagedAttention 将虚拟内存的理念应用于 attention 计算：

核心思想：不再将 KV cache 连续存储，而是将其拆分为固定大小的 block。每个请求获得一个 block table，将逻辑位置映射到 physical block。

虚拟内存类比

虚拟内存（操作系统） PagedAttention（vLLM）
页面（Page） block（如 16 个 token）
页表（Page Table） block table
物理内存 GPU 内存
内存分配器 block pool
缺页错误（Page Fault） cache miss
写时复制（Copy-on-Write） prefix 共享

基于 block 的存储

取代连续分配的方式：

1
2
3
4
5
6
7
8
9
10
11
Request 1 (512 tokens, 32 blocks):
Blocks: [11, 3, 7, 12, ..., 89]
       Scattered in GPU memory

Request 2 (1800 tokens, 113 blocks):
Blocks: [1, 2, 4, 5, ..., 105]
       Can reuse freed blocks from other requests

Request 3 (200 tokens, 13 blocks):
Blocks: [6, 8, 9, ..., 18]
       Minimal waste (only last block partially filled)

优势：

  1. 无过度分配：按需分配 block
  2. 无碎片化：任何空闲 block 都可分配给任意请求
  3. 近零浪费：每个请求只有最后一个 block 可能是部分填充的
  4. 共享：共同前缀可共享 physical block

实现深度解析

block 结构

文件位置：vllm/v1/core/kv_cache_utils.py

当前源码中的 KVCacheBlock 是一个 @dataclass(slots=True)，核心字段包括：

1
2
3
4
5
6
7
8
@dataclass(slots=True)
class KVCacheBlock:
    block_id: int
    ref_cnt: int = 0
    _block_hash: BlockHashWithGroupId | None = None
    prev_free_block: KVCacheBlock | None = None
    next_free_block: KVCacheBlock | None = None
    is_null: bool = False

其中 block_hash 通过 property 暴露，实际携带的是“block hash + KV cache group id”的组合键，而不是一个简单的 int。

物理存储：

v1 不再假设单一固定的 torch.empty(..., float16) layout。实际实现会先根据 KVCacheConfig 分配原始缓存，再由不同 attention backend（如 FlashInfer、Triton、FlashAttention）重塑为各自需要的 KV layout。常见 shape 仍然围绕 num_blocks / 2(K,V) / block_size / num_kv_heads / head_size 展开，但具体维度顺序与 dtype 取决于 backend。

Block Pool

文件位置：vllm/v1/core/block_pool.py

BlockPool 维护所有 KV cache block、空闲队列与 prefix cache 索引。构造函数的核心形态如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class BlockPool:
    def __init__(
        self,
        num_gpu_blocks: int,
        enable_caching: bool,
        hash_block_size: int,
        enable_kv_cache_events: bool = False,
        metrics_collector: KVCacheMetricsCollector | None = None,
    ):
        self.blocks = [KVCacheBlock(idx) for idx in range(num_gpu_blocks)]
        self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)
        self.cached_block_hash_to_block = BlockHashToBlockMap()

        # block 0 被预留为特殊 null_block
        self.null_block = self.free_block_queue.popleft()
        self.null_block.is_null = True

这有两个容易忽略的细节：

  1. block_id=0 不是普通数据块，而是 null_block 占位块。
  2. cached_block_hash_to_block 不是简单的 dict[int, block]，而是支持 hash 冲突的 BlockHashToBlockMap。

分配与释放：

BlockPool 的核心接口是 get_new_blocks()、touch() 与 free_blocks()：

1
2
3
new_blocks = block_pool.get_new_blocks(num_blocks)
block_pool.touch(cached_blocks)            # prefix cache 命中时增加引用
block_pool.free_blocks(reversed(blocks))   # 释放时通常按逆序归还
  • get_new_blocks() 会从 free queue 批量取出 block，并在需要时调用 _maybe_evict_cached_block() 清理旧的缓存元数据。
  • touch() 用于 prefix cache 命中路径：若命中的 block 仍在 free queue 中，会先移出再增加 ref_cnt。
  • free_blocks() 按列表批量释放 block；调用方通常按逆序释放，以便让尾部 block 更容易被优先回收。

block table 与 KVCacheBlocks

当前 v1 更接近“按 KV cache group 组织 block 列表”的实现，而不是 [[7], [23], ...] 这种每个逻辑 block 再包一层列表的结构。

  • 请求侧使用 KVCacheBlocks.blocks: tuple[Sequence[KVCacheBlock], ...]
  • 调用 get_block_ids() 后，会得到按 KV cache group 分组的 block id 列表
  • kernel 侧最终消费的是按 group 构造的 block_table tensor

对单一 attention group 来说，可以把它理解为一个平坦的 physical block id 序列；有多个 group（如 attention + Mamba）时，外层再按 group 分组。

KVCacheManager

文件位置：vllm/v1/core/kv_cache_manager.py

KVCacheManager 的职责不只是“给 request 分配新 block”，而是把 prefix cache 命中、完整序列准入、sliding window 裁剪、speculative lookahead，以及 encoder-decoder 的 cross-attention 都纳入同一套 slot 管理逻辑。

当前更重要的接口是：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
computed_blocks, num_computed_tokens = kv_cache_manager.get_computed_blocks(request)

can_admit = kv_cache_manager.can_fit_full_sequence(
    request,
    num_new_computed_tokens=num_computed_tokens,
    new_computed_blocks=computed_blocks,
)

new_blocks = kv_cache_manager.allocate_slots(
    request,
    num_new_tokens,
    num_new_computed_tokens=num_computed_tokens,
    new_computed_blocks=computed_blocks,
    num_lookahead_tokens=num_lookahead_tokens,
    num_external_computed_tokens=num_external_computed_tokens,
    delay_cache_blocks=delay_cache_blocks,
    num_encoder_tokens=num_encoder_tokens,
)

这套接口有几个关键点：

  • get_computed_blocks() 与 allocate_slots() 是分开的两个阶段，scheduler 会先查 prefix cache，再决定是否准入。
  • allocate_slots() 返回 KVCacheBlocks | None，而不是 (blocks, num_cached_tokens)。
  • allocate_slots() 还会在真正分配前调用 remove_skipped_blocks()，把 sliding window 或 sparse attention 下已经不需要的 block 替换成 null_block。

Prefix Caching：跨请求共享

Prefix caching 仍然是 PagedAttention 最有价值的能力之一，但当前实现比“对 token tuple 做一次 Python hash”要复杂得多。

block 哈希

源码里的 block hash 是链式计算的：当前 block 的 hash 会把父 block 的 hash、当前 block token、以及 multimodal / LoRA / cache salt / prompt embed 等额外键一起纳入计算。返回值类型是 BlockHash（bytes 包装），不是简单的 int。

可以把它理解为：

1
block_hash = H(parent_block_hash, current_block_tokens, extra_keys)

这让 vLLM 可以：

  • 只对完整 block建立 prefix cache 索引
  • 区分同一 token 序列在不同 KV cache group 中的缓存项
  • 在多模态、LoRA 或 prompt embedding 参与时避免误共享

cache 查找

当前 KVCacheManager.get_computed_blocks() 会调用 coordinator / manager 路径来寻找“最长可命中的前缀”，而不是直接遍历 request.block_hashes 并手动递增 ref_cnt。

1
2
3
4
5
6
computed_blocks, num_new_computed_tokens = (
    self.coordinator.find_longest_cache_hit(
        request.block_hashes,
        max_cache_hit_length,
    )
)

几个细节很重要：

  • prefix cache 只会命中完整 block
  • 即使“整个 prompt 都命中”，也通常要重新计算最后一个 token 以得到 logits，因此 max_cache_hit_length 会被限制为 request.num_tokens - 1
  • 命中后的引用计数增加发生在 block_pool.touch(...) 路径，而不是在查找函数里直接修改

reference counting

reference counting 的基本含义没有变化：多个 request 可以共享同一个 prefix block，只有 ref_cnt == 0 时它才会回到 free queue，成为可复用或可淘汰对象。

cache 淘汰

当前实现不是“扫描 free queue 找第一个 ref_cnt == 0 的 cached block 并返回”的 evict_cached_block() 接口。更接近真实源码的过程是：

  1. get_new_blocks() 先从 free queue 取出候选 block
  2. 若该 block 仍带有 block_hash，则调用 _maybe_evict_cached_block(block) 把它从 cached_block_hash_to_block 中摘除
  3. 若需要显式按 block id 淘汰 prefix cache，则使用 evict_blocks(block_ids)

也就是说，淘汰是围绕“取出的 block 是否还挂着 cache 元数据”展开的。

使用 block 进行 attention 计算

Prefill 阶段

当前 v1 并没有统一公开的 paged_attention_prefill(...) Python API。prefill 由具体 attention backend 负责执行，例如 FlashInfer 路径会使用 BatchPrefillWithPagedKVCacheWrapper / BatchPrefillWithRaggedKVCacheWrapper，其他 backend 则会构造自己的 metadata 和 kernel 调用。

对外可以把它理解为两步：

  1. scheduler / attn metadata 根据 seq_lens 和 block tables 构建 batch 视图
  2. backend kernel 直接按 paged KV cache layout 完成 prefill attention

Decode 阶段

decode 阶段通常由具体 attention backend 负责。以 FlashInfer 路径为例，decode 侧会通过 BatchDecodeWithPagedKVCacheWrapper（或 TRT-LLM decode path）直接消费 paged KV cache。

因此，当前实现更准确的描述是：

  • prefill / decode 都是 backend-specific
  • block table 会在 metadata 构建阶段转成 tensor 形式
  • kernel 直接根据 paged KV layout 取 K/V，而不是先在 Python 层把 block 拼接成连续张量

多组 KV cache

vLLM 支持不同层使用不同 cache 规格，相关配置类型是 KVCacheGroupSpec / KVCacheConfig。

概念上可以理解为：

  • KVCacheConfig 持有 num_blocks、kv_cache_tensors 和 kv_cache_groups
  • 每个 KVCacheGroupSpec 绑定一组 layer 名称与对应的 kv_cache_spec
  • kv_cache_spec 可以是 FullAttentionSpec、SlidingWindowSpec、MambaSpec 等不同类型

因此，一个 request 的 block 记录也会天然按 group 分层，例如：

1
2
3
4
blocks = (
    [attn_block_0, attn_block_1, attn_block_2],
    [mamba_block_0, mamba_block_1, mamba_block_2],
)

这也是为什么 BlockHash 需要和 group id 绑定：同一段 token 在不同 KV cache group 中对应的是不同的物理缓存项。

性能影响

内存效率

使用 PagedAttention 之前（连续分配）：

1
2
3
10 requests × 2048 max length × 512 KB/token = 10.5 GB
Actual usage: 10 requests × 500 avg length = 2.5 GB
Efficiency: 24%

使用 PagedAttention 之后（block 分配）：

1
2
3
500 blocks allocated across all requests
Waste: ~10 blocks (partial last blocks)
Efficiency: 98%

结果：并发请求数提升 4 倍！

吞吐量提升

真实基准测试（Llama-3-8B，运行于 H100）：

指标 未使用 PagedAttention 使用 PagedAttention 提升幅度
并发请求数 12 64 5.3x
吞吐量（tok/s） 1,500 8,000 5.3x
内存占用 60 GB 60 GB 持平
延迟（TTFT） 45ms 42ms -7%

Prefix Caching 收益

使用共同系统提示词（500 个 token）时：

请求次数 未使用缓存 使用缓存 加速比
第 1 次 10ms（prefill） 10ms 1x
第 2 次 10ms 1ms 10x
第 100 次 10ms 1ms 10x

缓存命中率：对于带有系统提示词的聊天机器人，通常在 60-80% 之间。

进阶主题

滑动窗口 attention

对 Mistral 这类 sliding-window 模型，当前实现不会简单地“切片旧 block 然后逐个 free”。更接近真实源码的流程是：

  • 先根据 total_computed_tokens 计算哪些 block 已经落到窗口之外
  • 调用 remove_skipped_blocks() 把这些位置替换为 null_block
  • 再把不再需要的真实 block 批量归还给 block pool

这样既保留了逻辑 block table 的形状，又避免了窗口外 token 继续占用可复用的缓存块。

speculative decoding

在当前 v1 中，speculative decoding 主要通过 num_lookahead_tokens、slot allocation 和 scheduler 的 scheduled_spec_decode_tokens 协同完成。

更准确地说：

  • scheduler 会为 lookahead token 预留 slot
  • allocate_slots() 只会把最终确认的 token 提交到 prefix cache
  • 被拒绝的 draft token 会在 slot mapping / 输出更新路径中被屏蔽，而不是靠手工 ref_cnt += 1 / 回滚 block 列表来管理

实现中的注意事项

hash 碰撞

当前实现通过 BlockHashToBlockMap 处理哈希碰撞。对于同一个 hash，值可能是单个 KVCacheBlock，也可能退化成 {block_id: KVCacheBlock} 这样的结构，而不是简单的 hash -> list[block]。

不完整的 block

最后一个 block 仍然可能是不完整的，但当前实现不会在 KVCacheBlock 元数据里单独维护 num_tokens_in_last_block 这样的字段。vLLM 只对完整 block建立 prefix cache，并通过 seq_lens 等 batch metadata 在 attention 计算时屏蔽无效位置。

线程安全

线程安全方面，block_pool.py 并没有对外暴露 allocate/_allocate_unsafe 这类接口；调度路径更接近单一 owner 管理 block pool 状态，而不是通过显式锁 API 暴露给外部使用。

核心要点

  1. PagedAttention 将 KV cache 拆分为固定大小的 block，消除了碎片化和过度分配问题

  2. block table 将逻辑 block 映射到 physical block，类似于虚拟内存的页表

  3. prefix caching 跨请求共享 block，大幅减少了冗余计算

  4. reference counting 确保安全共享，防止过早释放内存

  5. 近零内存浪费在实践中实现了 4-5 倍的吞吐量提升

  6. attention 内核经过优化，可直接使用 block 索引存储进行计算

下一步

在第三部分，我们将探索 Scheduler——vLLM 的核心调度器，它负责决定处理哪些请求、何时进行抢占，以及如何通过 continuous batching 最大化 GPU 利用率。

参考资料

PagedAttention 是 vLLM 卓越性能的基础。在下一篇文章中，我们将看到 Scheduler 如何在此之上协调请求的执行。