AI 推理性能调优：KV Cache 优化与显存管理的工程实践

一、显存墙：为什么大模型推理总是"卡在显存不够"

大模型推理的性能瓶颈往往不是计算力（FLOPS），而是显存带宽与容量。以 Llama-3-8B 为例，模型权重占用约 16GB（FP16），推理时还需要额外的 KV Cache 存储注意力键值对。KV Cache 的大小与序列长度和批大小线性相关：当序列长度为 4096、批大小为 32 时，KV Cache 可能占用 8-12GB 显存，总显存需求超过 24GB，单卡 A100 也捉襟见肘。KV Cache 优化是突破显存墙、提升推理吞吐的关键手段。

二、KV Cache 的内存模型与优化路径

KV Cache 的显存占用公式为：2 × num_layers × batch_size × seq_len × head_dim × num_kv_heads × dtype_size。其中 2 代表 Key 和 Value 各一份。优化路径有三条：降低精度（FP16→INT8/INT4）、减少序列长度（滑动窗口）、减少 KV Head 数量（GQA/MQA）。

graph TD
    A[KV Cache 显存优化] --> B[精度压缩<br/>FP16 → INT8/INT4]
    A --> C[序列截断<br/>滑动窗口注意力]
    A --> D[结构优化<br/>GQA / MQA]

    B --> B1[量化 KV Cache<br/>显存节省 50-75%]
    B --> B2[精度损失<br/>需校准评估]

    C --> C1[固定窗口大小<br/>显存占用恒定]
    C --> C2[长上下文丢失<br/>需配合 Sink Token]

    D --> D1[减少 KV Head 数<br/>显存线性下降]
    D --> D2[注意力质量下降<br/>需评估下游任务影响]

    style B fill:#e1f5fe
    style C fill:#c8e6c9
    style D fill:#fff3e0

GQA（Grouped-Query Attention）和 MQA（Multi-Query Attention）是目前最有效的结构优化方案。标准 MHA 中每个注意力头都有独立的 KV 对，GQA 将多个 Query Head 共享一组 KV，MQA 则所有 Query Head 共享一组 KV。Llama-3-8B 使用 GQA（8 组 KV Head），相比标准 MHA（32 组 KV Head），KV Cache 显存减少 75%。

三、KV Cache 优化的工程实现

3.1 KV Cache 量化

import torch
import numpy as np
from typing import Tuple

class KVCacheQuantizer:
    """
    KV Cache 量化器：将 FP16 的 KV Cache 量化为 INT8
    使用逐通道对称量化，保留每通道的缩放因子用于反量化

    设计考量：量化 KV Cache 与量化模型权重不同——
    KV Cache 是动态生成的，缩放因子需要在运行时实时计算，
    而非离线校准。逐通道量化比逐张量量化精度更高，
    因为不同通道的数值范围差异较大
    """

    @staticmethod
    def quantize_int8(tensor: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        将 FP16 张量量化为 INT8
        返回：(量化后的 INT8 张量, 缩放因子)
        """
        # 逐通道计算缩放因子：取绝对值最大值
        # tensor shape: [batch, num_heads, seq_len, head_dim]
        scale = tensor.abs().amax(dim=-1, keepdim=True) / 127.0
        # 避免除零：缩放因子最小值设为 1e-8
        scale = scale.clamp(min=1e-8)

        # 量化：缩放后四舍五入到 INT8 范围
        quantized = (tensor / scale).round().clamp(-128, 127).to(torch.int8)

        return quantized, scale.squeeze(-1)

    @staticmethod
    def dequantize_int8(
        quantized: torch.Tensor, scale: torch.Tensor
    ) -> torch.Tensor:
        """将 INT8 张量反量化为 FP16"""
        # scale shape: [batch, num_heads, seq_len]
        # quantized shape: [batch, num_heads, seq_len, head_dim]
        return quantized.float() * scale.unsqueeze(-1)

class KVCacheManager:
    """
    KV Cache 管理器：管理 KV Cache 的分配、复用与驱逐

    设计考量：PagedAttention 是当前最先进的 KV Cache 管理方案，
    将 KV Cache 按固定大小的 Page 分配，避免预分配连续显存。
    此处实现简化版的 Page 管理，展示核心逻辑
    """

    def __init__(
        self,
        num_layers: int,
        num_kv_heads: int,
        head_dim: int,
        page_size: int = 16,
        max_pages: int = 1024,
    ):
        self.num_layers = num_layers
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.page_size = page_size
        self.max_pages = max_pages

        # 空闲页面池
        self._free_pages = list(range(max_pages))
        # 每个请求占用的页面映射
        self._request_pages: dict = {}

    def allocate(self, request_id: str, num_tokens: int) -> list:
        """
        为请求分配 KV Cache 页面
        返回分配的页面 ID 列表
        """
        num_pages_needed = (num_tokens + self.page_size - 1) // self.page_size

        if len(self._free_pages) < num_pages_needed:
            # 显存不足：尝试驱逐最早完成的请求
            self._evict_oldest()

        if len(self._free_pages) < num_pages_needed:
            raise MemoryError(
                f"KV Cache 显存不足：需要 {num_pages_needed} 页，"
                f"可用 {len(self._free_pages)} 页"
            )

        allocated = self._free_pages[:num_pages_needed]
        self._free_pages = self._free_pages[num_pages_needed:]
        self._request_pages[request_id] = allocated

        return allocated

    def release(self, request_id: str):
        """释放请求占用的 KV Cache 页面"""
        if request_id in self._request_pages:
            pages = self._request_pages.pop(request_id)
            self._free_pages.extend(pages)

    def _evict_oldest(self):
        """驱逐最早完成的请求，释放其 KV Cache 页面"""
        if self._request_pages:
            oldest_id = next(iter(self._request_pages))
            self.release(oldest_id)

    def memory_usage(self) -> dict:
        """返回当前显存使用统计"""
        used_pages = self.max_pages - len(self._free_pages)
        bytes_per_page = (
            2  # Key + Value
            * self.num_layers
            * self.num_kv_heads
            * self.page_size
            * self.head_dim
            * 2  # FP16 = 2 bytes
        )
        used_bytes = used_pages * bytes_per_page
        total_bytes = self.max_pages * bytes_per_page

        return {
            "used_pages": used_pages,
            "total_pages": self.max_pages,
            "utilization": used_pages / self.max_pages,
            "used_gb": used_bytes / (1024 ** 3),
            "total_gb": total_bytes / (1024 ** 3),
        }

3.2 滑动窗口注意力实现

import torch
import torch.nn.functional as F

class SlidingWindowAttention:
    """
    滑动窗口注意力：限制每个 Token 只关注最近的 W 个 Token
    KV Cache 只保留最近 W 个位置的键值对，显存占用恒定

    设计考量：滑动窗口会丢失窗口外的上下文信息。
    Sink Token 策略保留序列开头的几个 Token（"注意力汇"），
    防止模型丢失全局信息（如 System Prompt）
    """

    def __init__(
        self,
        window_size: int = 4096,
        num_sink_tokens: int = 4,
    ):
        self.window_size = window_size
        self.num_sink_tokens = num_sink_tokens

    def compute_attention(
        self,
        query: torch.Tensor,   # [batch, num_heads, seq_len, head_dim]
        key: torch.Tensor,     # [batch, num_kv_heads, seq_len, head_dim]
        value: torch.Tensor,   # [batch, num_kv_heads, seq_len, head_dim]
    ) -> torch.Tensor:
        """计算滑动窗口注意力"""
        seq_len = query.shape[2]

        # 构建注意力掩码：滑动窗口 + Sink Token
        mask = torch.zeros(seq_len, seq_len, dtype=torch.bool)
        for i in range(seq_len):
            # 滑动窗口：每个位置只能看到前 window_size 个位置
            window_start = max(0, i - self.window_size + 1)
            mask[i, window_start:i + 1] = True
            # Sink Token：所有位置都能看到序列开头的几个 Token
            if self.num_sink_tokens > 0:
                mask[i, :self.num_sink_tokens] = True

        # 应用掩码：将不可见位置的注意力分数设为负无穷
        # 支持 GQA：如果 num_kv_heads < num_heads，需要扩展 key/value
        num_heads = query.shape[1]
        num_kv_heads = key.shape[1]
        if num_kv_heads < num_heads:
            n_rep = num_heads // num_kv_heads
            key = key.unsqueeze(2).expand(-1, -1, n_rep, -1, -1).reshape(
                key.shape[0], num_heads, key.shape[2], key.shape[3]
            )
            value = value.unsqueeze(2).expand(-1, -1, n_rep, -1, -1).reshape(
                value.shape[0], num_heads, value.shape[2], value.shape[3]
            )

        # Scaled Dot-Product Attention
        scale = query.shape[-1] ** -0.5
        scores = torch.matmul(query, key.transpose(-2, -1)) * scale
        scores = scores.masked_fill(~mask.to(scores.device), float("-inf"))
        weights = F.softmax(scores, dim=-1)
        output = torch.matmul(weights, value)

        return output

四、KV Cache 优化的边界与权衡

KV Cache 量化的精度损失是最大的隐忧。INT8 量化在大多数任务上的精度下降小于 1%，但在需要精细数值区分的任务（如数学推理、代码生成）上，精度下降可能达到 3-5%。INT4 量化的精度损失更显著，通常只在吞吐优先、精度容忍度高的场景（如对话补全）中使用。量化前必须在目标任务的基准测试集上评估精度影响。

滑动窗口注意力在长文本任务上存在信息丢失风险。窗口外的上下文被完全截断，模型无法"回忆"窗口外的内容。Sink Token 策略部分缓解了这个问题，但 Sink Token 数量有限，无法承载所有全局信息。对于需要全局上下文理解的任务（如文档摘要、长代码理解），滑动窗口不是合适的选择。

PagedAttention 的碎片化问题也需要关注。当请求的序列长度不是 Page 大小的整数倍时，最后一个 Page 会有空间浪费。Page 大小越小，碎片越少，但页面管理开销越大。生产环境通常选择 16-64 Token 的 Page 大小，在碎片率与管理开销之间取平衡。

五、总结

KV Cache 优化是突破大模型推理显存墙的核心手段。三条优化路径各有适用场景：精度压缩（INT8/INT4）适合吞吐优先场景，需评估精度损失；滑动窗口注意力适合短上下文对话场景，长文本任务需谨慎；GQA/MQA 是最有效的结构优化，已被主流模型采用。PagedAttention 解决了 KV Cache 的显存碎片问题，是当前生产环境的标准方案。优化选型应基于模型架构、任务特性和硬件配置综合决策。