大模型推理的 KV Cache 优化：从显存管理到前缀缓存

一、推理阶段的"显存墙"：KV Cache 的资源瓶颈

大模型推理时，每个 Token 的 Key 和 Value 向量需要缓存以供后续 Token 的注意力计算使用。这个 KV Cache 是自回归推理的核心数据结构——没有它，每生成一个 Token 都需要重新计算所有前序 Token 的 KV 向量，计算量随序列长度线性增长。

但 KV Cache 的显存占用极其惊人。以 LLaMA-2-70B 为例，FP16 精度下每个 Token 的 KV Cache 约 2.5MB，2048 Token 序列需要约 5GB 显存，仅 KV Cache 就占用了推理 GPU 的大量显存。当批量大小增加或序列长度扩展时，KV Cache 成为显存瓶颈，限制了推理吞吐量。

二、KV Cache 的底层机制与优化路径

KV Cache 的显存占用公式：显存 = 2 × num_layers × num_heads × head_dim × seq_len × batch_size × dtype_size。优化路径分为三类：压缩（减少每个元素的存储大小）、共享（复用不同请求间的 KV Cache）、回收（及时释放不再需要的缓存）。

flowchart TD
    A[KV Cache 显存优化] --> B[压缩: 减少存储大小]
    A --> C[共享: 复用缓存]
    A --> D[回收: 及时释放]
    B --> B1[量化: FP16→INT8/INT4]
    B --> B2[分页: PagedAttention 虚拟内存]
    C --> C1[前缀缓存: 系统 Prompt 复用]
    C --> C2[跨请求共享: 相同上下文]
    D --> D1[滑动窗口: 只保留最近 K 层]
    D --> D2[动态释放: 序列结束即回收]

三、KV Cache 优化的代码实现

3.1 PagedAttention：KV Cache 的虚拟内存管理

import torch
from dataclasses import dataclass
from typing import List, Optional

@dataclass
class BlockTable:
    """
    PagedAttention 的块表：将 KV Cache 划分为固定大小的块
    类似操作系统的虚拟内存分页机制
    每个序列的 KV Cache 由若干离散的块组成，无需连续显存
    """
    block_size: int = 16        # 每个块包含 16 个 Token 的 KV
    num_blocks: int = 1024      # 总块数
    block_tables: torch.Tensor  # [batch_size, max_num_blocks_per_seq]
    free_blocks: List[int]      # 空闲块列表

    def allocate_block(self) -> int:
        """分配一个空闲块，返回块 ID"""
        if not self.free_blocks:
            raise RuntimeError("KV Cache 显存不足: 无可用块")
        return self.free_blocks.pop()

    def free_block(self, block_id: int):
        """释放一个块，归还到空闲列表"""
        self.free_blocks.append(block_id)


class PagedKVCache:
    """
    PagedAttention KV Cache 管理器
    核心思路：将 KV Cache 按块管理，消除显存碎片
    不同序列可以共享相同的块（前缀缓存场景）
    """

    def __init__(self, num_layers: int, num_heads: int, head_dim: int,
                 block_size: int = 16, num_blocks: int = 4096,
                 dtype: torch.dtype = torch.float16):
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.block_size = block_size
        self.dtype = dtype

        # 预分配 KV Cache 显存池：所有块共享
        # 形状: [num_blocks, num_layers, 2, block_size, num_heads, head_dim]
        self.kv_buffer = torch.zeros(
            num_blocks, num_layers, 2, block_size, num_heads, head_dim,
            dtype=dtype, device='cuda'
        )

        # 块表管理
        self.free_blocks = list(range(num_blocks))
        self.sequence_blocks: dict[int, List[int]] = {}

    def allocate_sequence(self, seq_id: int, num_tokens: int) -> List[int]:
        """
        为新序列分配 KV Cache 块
        返回分配的块 ID 列表
        """
        num_needed = (num_tokens + self.block_size - 1) // self.block_size
        if len(self.free_blocks) < num_needed:
            raise RuntimeError(
                f"显存不足: 需要 {num_needed} 块, 可用 {len(self.free_blocks)} 块"
            )

        blocks = [self.free_blocks.pop() for _ in range(num_needed)]
        self.sequence_blocks[seq_id] = blocks
        return blocks

    def write_kv(self, seq_id: int, layer_idx: int,
                 positions: List[int], key: torch.Tensor, value: torch.Tensor):
        """
        将 KV 向量写入对应块的指定位置
        key/value: [num_tokens, num_heads, head_dim]
        """
        blocks = self.sequence_blocks[seq_id]
        for i, pos in enumerate(positions):
            block_idx = pos // self.block_size
            offset = pos % self.block_size
            block_id = blocks[block_idx]
            self.kv_buffer[block_id, layer_idx, 0, offset] = key[i]
            self.kv_buffer[block_id, layer_idx, 1, offset] = value[i]

    def free_sequence(self, seq_id: int):
        """序列结束后释放所有块"""
        if seq_id in self.sequence_blocks:
            self.free_blocks.extend(self.sequence_blocks.pop(seq_id))

3.2 前缀缓存：系统 Prompt 的 KV Cache 复用

class PrefixCacheManager:
    """
    前缀缓存管理器：复用系统 Prompt 的 KV Cache
    当多个请求共享相同的系统 Prompt 时，只需计算一次 KV
    后续请求直接复用，节省重复计算和显存
    """

    def __init__(self, kv_cache: PagedKVCache):
        self.kv_cache = kv_cache
        # 前缀哈希 → 块 ID 列表 + 引用计数
        self.prefix_cache: dict[str, tuple[List[int], int]] = {}

    def get_or_compute_prefix(self, prefix_text: str,
                               prefix_tokens: List[int],
                               model: torch.nn.Module) -> List[int]:
        """
        获取或计算前缀的 KV Cache
        如果缓存命中，直接返回块 ID 列表（引用计数 +1）
        如果缓存未命中，计算并缓存
        """
        prefix_hash = self._hash_prefix(prefix_tokens)

        if prefix_hash in self.prefix_cache:
            # 缓存命中：增加引用计数
            blocks, ref_count = self.prefix_cache[prefix_hash]
            self.prefix_cache[prefix_hash] = (blocks, ref_count + 1)
            return blocks

        # 缓存未命中：计算前缀 KV 并缓存
        with torch.no_grad():
            input_ids = torch.tensor([prefix_tokens], device='cuda')
            # 执行前向传播，只计算 KV，不生成 Token
            kv_data = model.compute_kv_cache(input_ids)

        # 分配块并写入 KV 数据
        blocks = self.kv_cache.allocate_sequence(
            seq_id=hash(prefix_hash),
            num_tokens=len(prefix_tokens)
        )
        for layer_idx in range(self.kv_cache.num_layers):
            self.kv_cache.write_kv(
                hash(prefix_hash), layer_idx,
                list(range(len(prefix_tokens))),
                kv_data[layer_idx][0], kv_data[layer_idx][1]
            )

        self.prefix_cache[prefix_hash] = (blocks, 1)
        return blocks

    def release_prefix(self, prefix_text: str, prefix_tokens: List[int]):
        """释放前缀缓存引用：引用计数归零时回收块"""
        prefix_hash = self._hash_prefix(prefix_tokens)
        if prefix_hash in self.prefix_cache:
            blocks, ref_count = self.prefix_cache[prefix_hash]
            if ref_count <= 1:
                # 引用计数归零：回收块
                for block_id in blocks:
                    self.kv_cache.free_blocks.append(block_id)
                del self.prefix_cache[prefix_hash]
            else:
                self.prefix_cache[prefix_hash] = (blocks, ref_count - 1)

    @staticmethod
    def _hash_prefix(tokens: List[int]) -> str:
        """计算 Token 序列的哈希值，用于缓存 Key"""
        import hashlib
        token_str = ','.join(map(str, tokens))
        return hashlib.sha256(token_str.encode()).hexdigest()[:16]

3.3 KV Cache 量化

class QuantizedKVCache:
    """
    KV Cache 量化：将 FP16 的 KV 向量量化为 INT8
    显存占用减半，精度损失在可接受范围内
    """

    def __init__(self, scale: float = 1.0):
        self.scale = scale

    def quantize(self, kv_tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        将 FP16 KV 量化为 INT8
        返回: (量化后的 INT8 张量, 缩放因子)
        每个注意力头独立计算缩放因子
        """
        # 计算每个头的最大绝对值
        max_val = kv_tensor.abs().amax(dim=-1, keepdim=True).clamp(min=1e-5)
        scale = max_val / 127.0

        # 量化: FP16 → INT8
        quantized = (kv_tensor / scale).round().clamp(-128, 127).to(torch.int8)
        return quantized, scale.squeeze(-1)

    def dequantize(self, quantized: torch.Tensor,
                   scale: torch.Tensor) -> torch.Tensor:
        """
        反量化: INT8 → FP16
        在注意力计算前执行
        """
        return quantized.float() * scale.unsqueeze(-1)

四、KV Cache 优化的边界分析与架构权衡

PagedAttention 的块碎片。块大小为 16 时，序列长度不是 16 的倍数时，最后一个块会有空闲空间。平均浪费率为 block_size / 2，即 8 个 Token 的空间。块大小越小浪费越少，但块表管理开销越大。

前缀缓存的一致性。系统 Prompt 更新后，旧的前缀缓存必须失效。如果使用版本号或哈希值作为缓存 Key，可以自动失效。但如果 Prompt 中包含动态内容（如当前日期），缓存命中率会大幅降低。

KV 量化的精度损失。INT8 量化在大多数任务上的精度损失小于 0.1%，但在需要高精度数值的任务（如数学推理）上，量化可能导致累积误差。INT4 量化的精度损失更显著，需要配合 GPTQ 等训练后量化方法。

适用边界：KV Cache 优化最适合高并发推理服务（如 ChatGPT 类应用）。对于单请求低延迟场景，PagedAttention 的块管理开销可能抵消显存节省的收益。

五、总结

KV Cache 是大模型推理的显存瓶颈，优化路径包括 PagedAttention 分页管理、前缀缓存复用和 KV 量化压缩。PagedAttention 消除了显存碎片，前缀缓存减少了重复计算，量化压缩降低了存储开销。落地时需关注块碎片浪费、缓存一致性和量化精度损失。建议从 PagedAttention + 前缀缓存开始，根据显存压力决定是否启用量化。