LLM 推理引擎内部原理:PagedAttention、Continuous Batching 与现代推理架构剖析

引言:推理才是 AI 的真正战场

当我们谈论大语言模型时,公众的目光总是聚焦在训练阶段——数千张 GPU 并行运算数周、耗费数百万美元。然而在产品落地中,推理(Inference)才是决定用户体验和运营成本的关键环节。一个 70B 参数模型在单次推理时,仅 KV Cache 就可能占用数十 GB 显存,而请求之间的序列长度差异会导致大量内存碎片和浪费。

vLLM 的出现改变了一切。它不是又一个推理框架,而是一场关于"如何像操作系统管理物理内存一样管理 KV Cache"的思维革命。本文将深入剖析现代 LLM 推理引擎的核心机制:Continuous Batching 调度策略和 PagedAttention 内存管理,并展示它们如何协同将推理吞吐量提升 2-4 倍。

一、问题本质:为什么朴素推理如此低效

1.1 自回归生成的本质瓶颈

LLM 是自回归模型:每次生成一个 token,然后将该 token 拼接到输入序列末尾,再次执行前向传播。这意味着推理过程中的计算图是动态增长的。

更关键的是,每个 token 的生成都需要访问之前所有 token 的 Key 和 Value 向量——这就是 KV Cache。对于长度为 S、hidden 维度为 D、层数为 L 的模型,单个请求的 KV Cache 大小为:

KV Cache 大小 = 2 × L × S × D × sizeof(dtype)

以 Llama 2 70B 为例(L=80, D=128, float16),一个 2048 token 的序列,KV Cache 约为:

2 × 80 × 2048 × 128 × 2 bytes ≈ 80 MB

而在实际服务场景中,成百上千个并发请求各自维护独立 KV Cache,显存压力呈爆炸式增长。

1.2 传统静态批处理的致命缺陷

最早期的推理服务采用静态批处理(Static Batching):将多个请求组成一个 batch,在前向传播中并行处理。这种方式简单直接,但存在一个致命问题——木桶效应:

请求 A: [token1, token2, token3]           ← 已生成完毕,等待...
请求 B: [token1, token2, token3, token4]   ← 还需要 1 步
请求 C: [token1, token2, ..., token20]     ← 还需要 17 步

整个 batch 必须等最长的请求完成才能释放资源。如果 batch 内有一个请求输出 2048 个 token,而其他请求早已结束,GPU 仍然在空转等待。据统计,静态批处理在真实负载下有 60%-80% 的计算浪费在处理 padding 和等待上。

1.3 连续批处理的直觉

连续批处理(Continuous Batching,最早由 DeepSpeed 提出,后被 vLLM 采用)的核心思想极为优雅:不等最慢的,快的先走。

在每个 decoding step 结束时,推理引擎检查当前 batch 中的请求状态。已完成的请求立即移出 batch 并释放资源,新到达的请求在下一个 step 立即加入。这意味着 GPU 的每个 decoding step 都在处理尽可能多的有效 token,几乎没有空转。

二、PagedAttention:用操作系统思维管理 KV Cache

2.1 显存碎片的根源

传统推理引擎为每个请求的 KV Cache 预分配一块连续的显存空间。问题在于:

  1. 长度未知:生成结束前,无法预知输出序列的最终长度
  2. 过度预留:必须按最大长度(如 2048)预分配,而大部分请求远不会用满
  3. 外部碎片:不同请求完成时间不同,释放的空间形成无法利用的碎片

vLLM 借鉴了操作系统的虚拟内存机制来解决这个问题。

2.2 从虚拟内存到 PagedAttention

操作系统中,进程看到的"连续内存"实际上是物理上不连续的页(Page)通过页表映射而成。PagedAttention 将同样的思想应用到 KV Cache 管理:

┌──────────────────────────────────────────────────────────┐
│                   PagedAttention 架构                      │
├──────────────────────────────────────────────────────────┤
│                                                          │
│  请求 A 的逻辑视图:  [KV0][KV1][KV2][KV3][KV4][KV5]...   │
│                          ↓ 页表映射                       │
│  物理 Block 分配:     Block3 → Block7 → Block1 → Block9  │
│                       (显存中任意位置,通过指针链接)         │
│                                                          │
│  请求 B 的逻辑视图:  [KV0][KV1][KV2]                      │
│                          ↓                                │
│  物理 Block 分配:     Block2 → Block5                     │
│                                                          │
└──────────────────────────────────────────────────────────┘

具体来说,PagedAttention 将每个请求的 KV Cache 分割成固定大小的 Block(通常 16 个 token 的 KV 向量 = 1 Block)。每个 Block 可以存放在显存的任意位置,通过一个 Block Table 维护逻辑顺序到物理位置的映射。

类比操作系统: - Block = 物理页帧(Page Frame) - Block Table = 页表(Page Table) - 请求的 KV 序列 = 进程的虚拟地址空间

2.3 Block 分配器:推理引擎的"内存管理单元"

vLLM 的 Block Engine 承担着类似 OS 中 buddy system 的角色:

class BlockEngine:
    def __init__(self, num_gpu_blocks: int, block_size: int = 16):
        self.block_size = block_size          # 每个 Block 存储的 token 数
        self.free_blocks = deque(range(num_gpu_blocks))  # 空闲 Block 池
        self.block_tables: dict[str, list[int]] = {}     # 请求ID → Block列表

    def allocate(self, request_id: str, num_tokens: int) -> list[int]:
        """为新请求/扩展请求分配 Block"""
        num_blocks = ceil(num_tokens / self.block_size)
        blocks = [self.free_blocks.popleft() for _ in range(num_blocks)]
        self.block_tables[request_id] = blocks
        return blocks

    def append_slot(self, request_id: str) -> Optional[int]:
        """生成新 token 时追加一个槽位"""
        blocks = self.block_tables[request_id]
        current_len = self.get_num_tokens(request_id)

        # 如果最后一个 Block 还有空间,直接写入
        if current_len % self.block_size != 0:
            return blocks[-1] * self.block_size + (current_len % self.block_size)

        # 否则分配新 Block
        if not self.free_blocks:
            return None  # OOM,需要抢占
        new_block = self.free_blocks.popleft()
        blocks.append(new_block)
        return new_block * self.block_size

    def free(self, request_id: str):
        """请求完成时释放所有 Block"""
        for block in self.block_tables.pop(request_id, []):
            self.free_blocks.append(block)

这个设计的精妙之处在于:一个请求的各个 Block 在物理上完全不需要连续。生成一个 token 时,如果当前 Block 已满,就分配一个空闲 Block(可能在显存的任意位置),通过 Block Table 维护顺序。这彻底消除了外部碎片。

2.4 CUDA Kernel 层面:如何高效访问不连续的 KV Cache

PagedAttention 的核心挑战在于:Attention 计算需要随机访问不同 Block 中的 K/V 向量。为此,vLLM 专门实现了分页 Attention CUDA kernel。

在标准 Attention 中,Q 与所有 K、V 做点积计算,K/V 在连续内存中。而 PagedAttention 需要在 kernel 中根据 Block Table 间接寻址:

// PagedAttention Kernel 伪代码
__global__ void paged_attention_kernel(
    float* out,           // 输出
    const float* q,       // 当前 query [num_heads, head_dim]
    const float* k_cache, // 全局 KV Cache 池 [num_blocks, block_size, num_heads, head_dim]
    const float* v_cache,
    const int* block_table, // 当前请求的 Block Table
    int block_table_len,
    int block_size,
    int num_heads,
    int head_dim
) {
    int head_idx = blockIdx.x;
    int tid = threadIdx.x;

    float acc = 0.0f;

    // 遍历该请求的所有 Block
    for (int block_idx = 0; block_idx < block_table_len; block_idx++) {
        int physical_block = block_table[block_table_len];  // 查表得到物理 Block 号

        // 在当前 Block 内遍历 token
        for (int token_offset = 0; token_offset < block_size; token_offset++) {
            // 计算 K 的物理地址
            int k_offset = physical_block * block_size * num_heads * head_dim
                          + token_offset * num_heads * head_dim
                          + head_idx * head_dim;

            // 执行 Q·K 点积
            float k_val = k_cache[k_offset + tid];
            float q_val = q[head_idx * head_dim + tid];
            acc += q_val * k_val;
        }
    }
    // ... 后续 softmax 和加权 V 的计算
}

这个 kernel 的核心创新在于两层间接寻址:先通过 Block Table 查到物理 Block 号,再在 Block 内按偏移访问。虽然引入了额外的内存间接访问,但由于 Block Size 通常设为 16,且现代 GPU 的 L2 Cache 能有效覆盖 Block Table 的访问热点,实际性能损失极小(相比传统实现差距在 3% 以内)。

三、完整推理循环:调度与执行的协同

3.1 vLLM 的调度器设计

vLLM 的调度器在每个 decoding step 执行一次,决定当前 step 处理哪些请求:

┌─────────────────────────────────────────────┐
│           调度器决策流程 (每个 step)           │
├─────────────────────────────────────────────┤
│                                             │
│  1. 将 RUNNING 请求加入当前 batch           │
│  2. 从 WAITING 队列填充剩余 slot            │
│  3. 检查显存是否足够                         │
│     └─ 不够 → 抢占(Preemption)             │
│  4. 组装 batch 并执行 forward               │
│  5. 检查完成的请求,释放资源                  │
│                                             │
└─────────────────────────────────────────────┘

调度器的核心逻辑如下:

class Scheduler:
    def schedule(self) -> SchedulerOutputs:
        # 阶段1: 尽可能多地放入 RUNNING 请求
        scheduled_seqs = []
        remaining_gpu_blocks = self.gpu_allocator.get_num_free_blocks()

        for seq_group in self.running:
            if remaining_gpu_blocks < 1:  # 至少需要一个新 Block
                break
            scheduled_seqs.append(seq_group)
            remaining_gpu_blocks -= 1  # 预留至少一个新 Block

        # 阶段2: 从 WAITING 队列填充
        for seq_group in list(self.waiting):
            num_blocks_needed = self._get_num_blocks(seq_group)
            if num_blocks_needed <= remaining_gpu_blocks:
                scheduled_seqs.append(seq_group)
                remaining_gpu_blocks -= num_blocks_needed
                self.waiting.remove(seq_group)
                self.running.append(seq_group)
            else:
                break

        # 阶段3: 显存不足时触发抢占
        if len(scheduled_seqs) == 0 and len(self.running) > 0:
            # 抢占最低优先级的 RUNNING 请求
            preempted = self._preempt_lowest_priority()
            scheduled_seqs.append(preempted)

        return SchedulerOutputs(
            scheduled_seq_groups=scheduled_seqs,
            blocks_to_swap_in={},
            blocks_to_swap_out={},
            blocks_to_copy={},  # 用于 Copy-on-Write
        )

3.2 抢占机制(Preemption):当显存耗尽时

当所有 Block 都被占用且新请求到达时,vLLM 必须通过抢占回收显存。有两种策略:

Recomputation(重算):直接释放被抢占请求的所有 Block,请求重新进入 WAITING 队列。恢复时从头重新计算 KV Cache。这种方式简单,但浪费之前的计算。

Swapping(换出):将 Block 内容从 GPU 显存换出到 CPU 内存,释放 GPU Block。恢复时再换回。这种方式更快,但需要 CPU-GPU 带宽。

vLLM 的实现策略是:当抢占请求的序列较短(Block 数量少)时选择 Recomputation(换入换出的传输开销比重算还大);当序列很长时选择 Swapping。

3.3 Copy-on-Write:并发生成的内存共享

PagedAttention 带来的一个重要优化是Copy-on-Write(CoW),用于并行采样场景。

当用户对同一个 prompt 请求多个采样(如 temperature sampling 想要 5 个不同回复)时,传统做法是为每个采样序列各复制一份完整 KV Cache。而 CoW 机制下,多个子序列共享同一段前缀 Block:

                    共享前缀 Block                分叉后的独立 Block

采样1: [Block0] [Block1] [Block2]  →  [Block5] [Block6]
采样2: [Block0] [Block1] [Block2]  →  [Block7] [Block8]
采样3: [Block0] [Block1] [Block2]  →  [Block9] [Block10]
                                      ↑
                             写时复制:只有当某个采样在
                             Block2 之后生成新 token 时,
                             才真正分配新 Block

当共享 Block 中的某个 token 需要修改(实际上 KV Cache 是追加写,不会原地修改新 Block),此时才触发复制。这种设计在 best_of > 1 或 beam_search 场景下,显存节省可达 (N-1)/N(N 为并发生成数)。

四、工程实践中的关键数值

4.1 显存预算计算

部署一个 7B 模型到单张 A100 80GB 上:

组件                      占用显存
───────────────────────────────────────
模型权重 (fp16, 7B)       ≈ 14 GB
CUDA 上下文 + 系统预留    ≈ 5 GB
                       ───────
可用 KV Cache 空间         ≈ 61 GB

每 token 每层的 KV 大小:
  2(key+value) × 32层 × 128(dim) × 2(fp16) ≈ 16 KB/token

可用 token 数 ≈ 61 GB / 16 KB ≈ 3.8M tokens (理论值)

典型场景: block_size=16 tokens/block
可分配 block 数 ≈ 3.8M / 16 ≈ 237,500 blocks

4.2 与 vLLM 交互的实践示例

from vLLM import LLM, SamplingParams

# 初始化引擎
llm = LLM(
    model="meta-llama/Llama-2-70b-chat-hf",
    tensor_parallel_size=4,        # 4 卡张量并行
    dtype="float16",
    max_num_seqs=256,              # 最大并发序列数
    max_num_batched_tokens=4096,   # 单 step 最大 token 数
    block_size=16,                 # PagedAttention Block 大小
    gpu_memory_utilization=0.9,    # 90% 显存用于 KV Cache
)

# 混合请求 - 这正是 PagedAttention 发挥优势的场景
prompts = [
    # 短请求:简单问题
    "用一句话解释什么是量子纠缠。",
    # 中等请求:代码生成
    "实现一个线程安全的 LRU 缓存,支持 get 和 put 操作,时间复杂度 O(1)。",
    # 长请求:文档摘要
    "请总结以下长文...(假设这里有一篇 5000 字的文章)",
]

outputs = llm.generate(
    prompts,
    SamplingParams(
        temperature=0.7,
        max_tokens=2048,
        top_p=0.9,
    )
)

# 在 Continuous Batching 下:
# - 短请求完成后立即释放 Block
# - 长请求继续使用新释放的 Block
# - 在任意 step,batch 中的请求集合可能不同

4.3 性能对比数据

在 vLLM 论文(Kwon et al., SOSP 2023)中报告的典型数据(Llama 1 13B, A100, 在线服务场景):

框架 吞吐量 (tokens/s) 显存利用率 平均 TTFT
HuggingFace Pipeline 基准 (1x) ~40% 较高
HuggingFace + Static Batching 1.5-2x ~50% 中等
vLLM (Continuous Batching + PagedAttention) 2-4x ~95% 最低

关键不是峰值吞吐,而是在混合长度负载下的稳定性。PagedAttention 消除了内存碎片,使得在持续高负载下不会出现初始阶段流畅但随着运行时间增长逐渐卡顿的"内存碎片墙"问题。

五、前沿演进:从 PagedAttention 到更多

5.1 分层的 KV Cache 管理(vLLM v0.4+)

最新版本的 vLLM 引入了前缀缓存(Prefix Caching):当多个请求共享相同的系统 prompt 前缀时,不必重复计算和缓存对应的 KV 向量。

请求1: [System Prompt: 你是一个有用的助手...] [User: 帮我写代码]
                                     ↑ 共享前缀,KV Block 复用
请求2: [System Prompt: 你是一个有用的助手...] [User: 帮我写文案]

5.2 推测解码与 PagedAttention 的协同

推测解码(Speculative Decoding)使用小模型(Draft Model)生成候选 token,再由大模型验证。PagedAttention 天然适配这种模式——候选 token 和验证 token 的 KV 可以共享 Block 结构,验证不通过的 token 对应的 Block 直接丢弃,无需考虑碎片整理。

5.3 异构显存分层

PagedAttention 的 Block 抽象使得"分级存储"成为可能:热 Block 留在 GPU 显存,冷 Block 迁移到 CPU 内存甚至 NVMe SSD。这种分层策略正在 RLHF 场景的大 batch 训练推理混合集群中发挥作用。

六、总结:操作系统思维在 AI 基础设施中的胜利

PagedAttention 给我们的启示远超于一个具体的模型推理优化。它揭示了一个更普遍的工程真相:当 GPU 成为新的"计算机",推理引擎成为新的"操作系统"时,经典的操作系统原理(虚拟内存、分页、调度、抢占)将以新的形态重生。

对于从事 AI 基础设施的工程师而言,理解 PagedAttention 不亚于当年理解 Linux 内存管理子系统——它让你从"显存占用过高"的直观抱怨,上升到对 Block 分配策略、抢占阈值、CoW 触发条件的系统级思考。

vLLM 的 PagedAttention 证明了:AI infra 的深层优化不在于堆更多 GPU,而在于如何像 OS 管理物理内存那样,优雅而高效地管理每一字节 KV Cache。


参考资料: - Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention", SOSP 2023 - vLLM 开源项目: https://github.com/vllm-project/vllm - DeepSpeed-FastGen: https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepseed-fastgen

点赞(0) 打赏

评论列表 共有 0 条评论

暂无评论
立即
投稿

微信公众账号

微信扫一扫加关注

发表
评论
返回
顶部