KV Cache内存优化技术原理与PagedAttention分页注意力实现

大模型推理过程中，KV Cache内存优化是提升吞吐量的关键环节。传统注意力机制在生成每个token时需要缓存此前所有token的Key和Value向量，随着序列长度增长，显存占用呈线性膨胀。PagedAttention通过借鉴操作系统的虚拟内存分页机制，将KV Cache划分为固定大小的Block，实现了显存碎片的大幅消除和并发批处理能力的显著提升。本文围绕KV Cache内存优化技术，拆解PagedAttention的分页存储原理并给出工程实现路径。

KV Cache内存模型与推理瓶颈分析

自回归生成模型在推理时，每一步都需要访问此前所有位置的Key和Value张量。对于Llama-2-7B这类模型，单个token的KV Cache在float16精度下约占0.5MB，2048 tokens的上下文窗口需要约1GB显存。当多个请求并发时，KV Cache的显存占用远超模型权重本身。

传统的内存分配方式为每个请求预分配一块连续的物理显存，其大小等于最大序列长度对应的KV Cache体积。这种方案存在两个核心问题：

– 内部碎片：实际生成长度通常远小于最大长度，预分配的显存大量浪费。实测数据显示，平均浪费率在60%到80%之间。

– 外部碎片：不同请求的生命周期不同，频繁分配和释放导致显存空间碎片化，新请求无法找到足够大的连续块。

PagedAttention分页存储机制详解

PagedAttention的核心思路是将KV Cache的物理存储拆分成固定大小的Block，每个Block默认包含16个token的KV数据。逻辑上连续的KV Cache通过Block Table映射到物理Block，类似操作系统的页表机制。

# PagedAttention Block Table 结构示意
# 每个请求维护一个Block Table，记录逻辑Block到物理Block的映射

class BlockTable:
    def __init__(self, block_size=16):
        self.block_size = block_size
        self.table = []  # logical_block_idx -> physical_block_idx

    def allocate(self, num_tokens, physical_pool):
        """分配物理Block并建立映射"""
        num_blocks = (num_tokens + self.block_size - 1) // self.block_size
        for _ in range(num_blocks):
            phys_block = physical_pool.get_free_block()
            self.table.append(phys_block)

    def get_physical_block(self, logical_idx):
        """逻辑Block索引转物理Block索引"""
        if logical_idx >= len(self.table):
            raise IndexError(f"逻辑Block {logical_idx} 超出范围")
        return self.table[logical_idx]

这种设计的直接收益在于：物理Block可以不连续，消除了外部碎片；Block按需分配，已满的Block才申请新Block，内部碎片从序列级别降低到Block级别（16 tokens），浪费率降至不到5%。

并发批处理中的内存复用策略

PagedAttention的另一个关键优势是支持灵活的批处理调度。传统方案要求一个Batch内所有请求同时到达、同时结束，而PagedAttention允许请求在不同时刻加入和退出Batch。

当某个请求生成完成后，其占用的物理Block立即归还到Free Pool，可供新请求复用。调度器可以动态调整Batch的组成，将等待中的请求填入刚释放的显存槽位。这种机制在vLLM中的实现称为Continuous Batching（连续批处理），实测可将吞吐量提升2到4倍。

# 连续批处理调度伪代码

class ContinuousBatchScheduler:
    def __init__(self, max_batch_size, block_pool, model):
        self.max_batch_size = max_batch_size
        self.block_pool = block_pool
        self.model = model
        self.running = []      # 正在运行的请求
        self.waiting = []      # 等待中的请求

    def schedule(self):
        # 移出已完成的请求，释放Block
        self.running = [r for r in self.running if not r.finished]
        for r in self.running:
            if r.finished:
                self.block_pool.free_blocks(r.block_table)

        # 将等待队列中的请求加入运行队列
        while (len(self.running) < self.max_batch_size
               and len(self.waiting) > 0):
            new_req = self.waiting.pop(0)
            try:
                new_req.block_table.allocate(
                    new_req.context_len, self.block_pool
                )
                self.running.append(new_req)
            except OutOfMemory:
                self.waiting.insert(0, new_req)
                break

        return self.running

KV Cache量化压缩与多精度存储方案

除了分页存储，量化是另一种有效的KV Cache优化手段。FP8量化将每个KV元素的精度从float16降到8位，显存占用减半，精度损失在多数任务上可接受。更激进的方案是INT4量化，将显存占用降至原来的四分之一。

实践中常采用混合精度策略：近期token使用float16保证精度，远期token降级到int8或int4节省空间。PagedAttention的Block粒度天然支持这种混合精度方案——不同Block可以指定不同的存储格式。

# 混合精度Block分配示例

MIXED_PRECISION_POLICY = {
    "recent_blocks": 32,      # 最近32个Block使用fp16
    "mid_blocks": "int8",     # 中间Block使用int8
    "old_blocks": "int4"      # 远期Block使用int4
}

def allocate_mixed_block(logical_idx, policy, pool):
    if logical_idx < policy["recent_blocks"]:
        return pool.get_block(dtype="float16")
    elif logical_idx < policy["recent_blocks"] * 4:
        return pool.get_block(dtype="int8")
    else:
        return pool.get_block(dtype="int4")

Prefix Sharing与共享Block优化技巧

在对话系统和Few-shot场景中，多个请求往往共享相同的前缀（System Prompt、Few-shot示例等）。PagedAttention支持Copy-on-Write的Block共享：多个请求的Block Table指向同一物理Block，当某个请求需要修改该Block时才分配新Block。

这一特性在vLLM中称为Automatic Prefix Caching（自动前缀缓存）。对于Chat场景，System Prompt部分可被数百个并发请求共享，节省的显存直接转化为更高的并发度。实测在Llama-2-13B模型上，开启前缀缓存后并发请求数提升约40%。

性能基准测试与调优参数推荐

以下是vLLM在不同配置下的吞吐量对比数据（Llama-2-7B，A100-80GB，2048上下文长度）：

调优建议：Block Size设为16适用于大多数场景，过小会增加Block Table的内存开销，过大会增加内部碎片。GPU Memory Utilization参数建议设为0.90，预留10%作为临时缓冲。Maximum Sequences参数应根据GPU显存和模型大小动态计算，公式为（可用显存乘以利用率减去模型权重）除以（单token KV Cache大小乘以平均序列长度）。