Chaos and Order

💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.

引言
Transformer Self-Attention 与 KV Cache 基础
- Self-Attention 运算
- KV Cache 内存分析
Multi-Head Attention（MHA）内存分析
- MHA 结构
Multi-Query Attention（MQA）
- MQA 结构与节省效果
Grouped Query Attention（GQA）
- GQA 架构（Llama 2/3）
Multi-Head Latent Attention（MLA）
- MLA 架构（DeepSeek-V2/V3）
MHA vs MQA vs GQA vs MLA 对比表
KV Cache 压缩技术
PagedAttention（vLLM）
- 基于虚拟内存的 KV Cache 管理
故障案例与恢复流程
优化清单
结语
参考资料

KV Cache Optimization: GQA, MLA, MHA Attention Mechanisms

引言

大规模语言模型（LLM）推理成本中最大的瓶颈，是 KV Cache（Key-Value Cache） 的内存消耗。当 GPT-4 级别的模型以 128K 上下文长度同时处理数百个请求时，仅 KV Cache 就可能消耗数百 GB 的 GPU 内存。这种内存限制，直接成为并发吞吐量（throughput）与响应延迟（latency）的制约因素。

在 Transformer 的 Self-Attention 机制中，每个解码步骤都重新计算此前所有 token 的 Key 和 Value 向量是低效的，因此把它们存入缓存并复用，正是 KV Cache 的基本原理。问题在于，序列长度越长，这个缓存的大小就线性增长。

为解决这个问题，业界提出了多种注意力机制。Multi-Query Attention（MQA）让所有 Query 头共享同一个 KV 头，从而大幅降低内存，但也带来了质量下降。Grouped Query Attention（GQA）作为 MHA 与 MQA 之间的折衷方案，被 Llama 2/3 采用。Multi-Head Latent Attention（MLA）在 DeepSeek-V2/V3 中把 KV 压缩到低维潜在空间，取得了惊人的效率。

本文将系统性地梳理各注意力机制的数学原理与内存分析、KV Cache 压缩技术、PagedAttention（vLLM）、PyTorch 实现示例、真实的 OOM 故障案例与恢复过程，以及优化清单。

Transformer Self-Attention 与 KV Cache 基础

Self-Attention 运算

Transformer 的 Scaled Dot-Product Attention 定义如下。

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

其中，Q（Query）、K（Key）、V（Value）是通过对输入 token 的嵌入做线性变换得到的。

import torch
import torch.nn as nn
import math

class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, kv_cache=None):
        B, T, C = x.shape
        # Q, K, V 投影
        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        # 利用 KV Cache
        if kv_cache is not None:
            k_cache, v_cache = kv_cache
            k = torch.cat([k_cache, k], dim=2)
            v = torch.cat([v_cache, v], dim=2)

        # Attention 运算
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn, v)

        output = output.transpose(1, 2).contiguous().view(B, T, C)
        return self.W_o(output), (k, v)

KV Cache 内存分析

KV Cache 的内存占用可用以下公式计算。

KV Cache 内存 = 2（K 和 V）x 层数（L）x 头数（H）x 序列长度（S）x 头维度（d_k）x 字节大小

例如，以 Llama 2 70B 模型为例：

层数：80
头数：64（GQA 之前的基准）
头维度：128
序列长度：4096
FP16（2 字节）

KV Cache = 2 x 80 x 64 x 4096 x 128 x 2 = 10.7 GB（单个请求）

若批大小为 32，则为 342.4 GB，远远超过模型权重本身的内存占用。

Multi-Head Attention（MHA）内存分析

MHA 结构

在标准的 Multi-Head Attention 中，每个注意力头都拥有独立的 Q、K、V 投影。H 个头各自维持 d_k 维的 K、V，因此 KV Cache 达到最大规模。

class MultiHeadAttention(nn.Module):
    """标准 Multi-Head Attention（MHA）
    每个头维持独立的 KV 对"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)  # H * d_k 维
        self.W_v = nn.Linear(d_model, d_model)  # H * d_k 维
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, kv_cache=None):
        B, T, _ = x.shape

        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        if kv_cache is not None:
            k = torch.cat([kv_cache[0], k], dim=2)
            v = torch.cat([kv_cache[1], v], dim=2)

        # KV Cache 大小: [B, n_heads, S, d_k] x 2
        # 内存: 2 * B * n_heads * S * d_k * sizeof(dtype)
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out), (k, v)

MHA 的 KV Cache per token per layer = 2 x H x d_k x sizeof(dtype)

Multi-Query Attention（MQA）

MQA 结构与节省效果

MQA 是 Shazeer（2019）提出的方法，所有 Query 头 共享单一的 KV 头。KV Cache 缩小 H 倍，但会观察到质量下降。

class MultiQueryAttention(nn.Module):
    """Multi-Query Attention（MQA）
    所有 Query 头共享单一 KV 头"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)       # H * d_k
        self.W_k = nn.Linear(d_model, self.d_k)       # 1 * d_k（单一头）
        self.W_v = nn.Linear(d_model, self.d_k)       # 1 * d_k（单一头）
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, kv_cache=None):
        B, T, _ = x.shape

        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, 1, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, 1, self.d_k).transpose(1, 2)

        if kv_cache is not None:
            k = torch.cat([kv_cache[0], k], dim=2)
            v = torch.cat([kv_cache[1], v], dim=2)

        # 把 K、V 广播到所有头
        k = k.expand(-1, self.n_heads, -1, -1)
        v = v.expand(-1, self.n_heads, -1, -1)

        # KV Cache 大小: [B, 1, S, d_k] x 2
        # 相比 MHA 缩小为 1/H
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out), (k[:, :1], v[:, :1])

MQA 的 KV Cache per token per layer = 2 x 1 x d_k x sizeof(dtype)，相比 MHA 缩小为 1/H。

Grouped Query Attention（GQA）

GQA 架构（Llama 2/3）

GQA 由 Ainslie 等（2023）提出，并被 Llama 2、Llama 3 采用，它把 Query 头分成 G 个组，每个组共享一个 KV 头。当 G=1 时等价于 MQA，G=H 时等价于 MHA。

class GroupedQueryAttention(nn.Module):
    """Grouped Query Attention（GQA）
    将 Query 头分成 G 个组，各组共享 KV
    被 Llama 2/3 采用"""
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        self.n_heads = n_heads        # Query 头数
        self.n_kv_heads = n_kv_heads  # KV 头数（组数）
        self.d_k = d_model // n_heads
        self.n_rep = n_heads // n_kv_heads  # 每组的 Query 头数

        self.W_q = nn.Linear(d_model, n_heads * self.d_k)
        self.W_k = nn.Linear(d_model, n_kv_heads * self.d_k)
        self.W_v = nn.Linear(d_model, n_kv_heads * self.d_k)
        self.W_o = nn.Linear(d_model, d_model)

    def repeat_kv(self, x):
        """把 KV 头重复到与 Query 头数一致"""
        B, H_kv, S, D = x.shape
        if self.n_rep == 1:
            return x
        return (
            x[:, :, None, :, :]
            .expand(B, H_kv, self.n_rep, S, D)
            .reshape(B, self.n_heads, S, D)
        )

    def forward(self, x, kv_cache=None):
        B, T, _ = x.shape

        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_kv_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_kv_heads, self.d_k).transpose(1, 2)

        if kv_cache is not None:
            k = torch.cat([kv_cache[0], k], dim=2)
            v = torch.cat([kv_cache[1], v], dim=2)

        # KV Cache 大小: [B, n_kv_heads, S, d_k] x 2
        new_cache = (k, v)

        # 为了广播而重复 KV 头
        k = self.repeat_kv(k)
        v = self.repeat_kv(v)

        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out), new_cache

Llama 3 70B 的 GQA 设置：n_heads=64，n_kv_heads=8。KV Cache 相比 MHA 缩小为 1/8，而质量几乎没有下降。

Multi-Head Latent Attention（MLA）

MLA 架构（DeepSeek-V2/V3）

MLA 是 DeepSeek-V2（2024）提出的创新方法，把 KV 向量压缩为低维的 潜在向量（latent vector）存入缓存。解码时再从潜在向量中还原出 K、V。

核心思路如下：

把输入 x 压缩为低维潜在向量 c：c = W_down * x
从潜在向量还原出 K、V：K = W_uk _ c，V = W_uv _ c
缓存中只保存低维的 c

class MultiHeadLatentAttention(nn.Module):
    """Multi-Head Latent Attention（MLA）
    用于 DeepSeek-V2/V3
    将 KV 压缩到低维潜在空间后再缓存"""
    def __init__(self, d_model, n_heads, d_latent, rope_dim=64):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.d_latent = d_latent  # 潜在维度（远小于 d_model）
        self.rope_dim = rope_dim

        # Query 投影
        self.W_q = nn.Linear(d_model, n_heads * self.d_k)

        # KV 压缩（Down-projection）
        self.W_dkv = nn.Linear(d_model, d_latent)

        # KV 还原（Up-projection）
        self.W_uk = nn.Linear(d_latent, n_heads * self.d_k)
        self.W_uv = nn.Linear(d_latent, n_heads * self.d_k)

        # 用于 RoPE 的单独投影
        self.W_kr = nn.Linear(d_model, rope_dim)

        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, kv_cache=None):
        B, T, _ = x.shape

        # Query
        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        # KV 压缩: d_model -> d_latent
        c_kv = self.W_dkv(x)  # [B, T, d_latent]

        # RoPE 键
        k_rope = self.W_kr(x)  # [B, T, rope_dim]

        if kv_cache is not None:
            c_kv_cached, k_rope_cached = kv_cache
            c_kv = torch.cat([c_kv_cached, c_kv], dim=1)
            k_rope = torch.cat([k_rope_cached, k_rope], dim=1)

        # KV Cache: 只保存低维的 c_kv 和 k_rope
        # 内存: B * S * (d_latent + rope_dim) * sizeof(dtype)
        # 相比 MHA: 按 d_latent / (2 * n_heads * d_k) 的比例缩小
        new_cache = (c_kv, k_rope)

        # KV 还原: d_latent -> n_heads * d_k
        S = c_kv.shape[1]
        k = self.W_uk(c_kv).view(B, S, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_uv(c_kv).view(B, S, self.n_heads, self.d_k).transpose(1, 2)

        # Attention 运算
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out), new_cache

DeepSeek-V2 的 MLA 设置：d_model=5120，d_latent=512，n_heads=128。KV Cache 相比 MHA 缩小了约 93.7%，同时达到了与 MHA 相当的性能。

MHA vs MQA vs GQA vs MLA 对比表

项目	MHA	MQA	GQA	MLA
KV 头数	H	1	G (1 < G < H)	-（潜在向量）
每 token 每层缓存	2Hd_k	2d_k	2Gd_k	d_latent + d_rope
相对 MHA 的缓存比例	100%	1/H	G/H	d_latent/(2Hd_k)
Llama 3 70B 示例（H=64, G=8）	16,384B	256B	2,048B	-
DeepSeek-V2 示例	163,840B	-	-	约 4,608B
质量影响	基线	略有下降	几乎相同	几乎相同
推理速度	基线	最快	良好	良好（有还原成本）
训练稳定性	高	中等	高	高
代表模型	GPT-3, BERT	PaLM, Falcon	Llama 2/3, Mistral	DeepSeek-V2/V3
实现复杂度	低	低	中	高

KV Cache 压缩技术

量化（Quantization）

把 KV Cache 从 FP16 量化为 INT8 或 INT4，可以把内存减少 2-4 倍。

class QuantizedKVCache:
    """KV Cache 量化
    FP16 -> INT8 转换，节省 50% 内存"""
    def __init__(self, n_layers, n_heads, max_seq_len, d_k, dtype=torch.int8):
        self.n_layers = n_layers
        self.scales = {}  # 量化缩放因子
        self.zero_points = {}
        self.cache_k = {}
        self.cache_v = {}

    def quantize(self, tensor):
        """按通道（Per-channel）做 INT8 量化"""
        min_val = tensor.min(dim=-1, keepdim=True)[0]
        max_val = tensor.max(dim=-1, keepdim=True)[0]
        scale = (max_val - min_val) / 255.0
        zero_point = (-min_val / scale).round().clamp(0, 255)
        quantized = ((tensor / scale) + zero_point).round().clamp(0, 255).to(torch.uint8)
        return quantized, scale, zero_point

    def dequantize(self, quantized, scale, zero_point):
        """INT8 -> FP16 反量化"""
        return (quantized.float() - zero_point) * scale

    def update(self, layer_idx, k, v):
        q_k, s_k, z_k = self.quantize(k)
        q_v, s_v, z_v = self.quantize(v)
        if layer_idx in self.cache_k:
            self.cache_k[layer_idx] = torch.cat([self.cache_k[layer_idx], q_k], dim=2)
            self.cache_v[layer_idx] = torch.cat([self.cache_v[layer_idx], q_v], dim=2)
        else:
            self.cache_k[layer_idx] = q_k
            self.cache_v[layer_idx] = q_v
        self.scales[layer_idx] = (s_k, s_v)
        self.zero_points[layer_idx] = (z_k, z_v)

    def get(self, layer_idx):
        k = self.dequantize(
            self.cache_k[layer_idx],
            self.scales[layer_idx][0],
            self.zero_points[layer_idx][0]
        )
        v = self.dequantize(
            self.cache_v[layer_idx],
            self.scales[layer_idx][1],
            self.zero_points[layer_idx][1]
        )
        return k, v

驱逐策略（Eviction Policy）

随着序列变长，选择性地清除旧 token 的 KV 的策略。

H2O（Heavy-Hitter Oracle）方式只保留注意力分数较高的 token 的 KV。

class H2OKVCache:
    """Heavy-Hitter Oracle（H2O）KV Cache 驱逐
    只保留注意力分数较高的 token"""
    def __init__(self, max_cache_size, n_heads, d_k):
        self.max_cache_size = max_cache_size
        self.attention_scores = None

    def update(self, k, v, attn_weights):
        """基于注意力权重累积重要度并执行驱逐"""
        B, H, S_new, _ = k.shape

        if self.attention_scores is None:
            self.attention_scores = attn_weights.sum(dim=2)  # [B, H, S]
        else:
            # 累积新 token 的注意力分数
            new_scores = attn_weights.sum(dim=2)
            self.attention_scores = torch.cat(
                [self.attention_scores, new_scores[:, :, -S_new:]], dim=2
            )

        # 缓存超出大小时执行驱逐
        current_size = self.attention_scores.shape[2]
        if current_size > self.max_cache_size:
            # 识别重要度较低的 token（第一个 token 始终保留）
            scores = self.attention_scores[:, :, 1:]  # 排除第一个 token
            _, indices = scores.topk(
                self.max_cache_size - 1, dim=2, sorted=False
            )
            indices = indices + 1  # 偏移量修正
            # 加入第一个 token 的索引
            first_token = torch.zeros(B, H, 1, dtype=torch.long, device=k.device)
            keep_indices = torch.cat([first_token, indices], dim=2)

            # 只保留选中的 token
            k = k.gather(2, keep_indices.unsqueeze(-1).expand(-1, -1, -1, k.shape[-1]))
            v = v.gather(2, keep_indices.unsqueeze(-1).expand(-1, -1, -1, v.shape[-1]))
            self.attention_scores = self.attention_scores.gather(2, keep_indices)

        return k, v

Sliding Window Attention

Mistral 等模型中使用的固定大小窗口注意力。

class SlidingWindowAttention(nn.Module):
    """Sliding Window Attention
    只对最近的 W 个 token 应用注意力
    在 Mistral、Gemma 等模型中使用"""
    def __init__(self, d_model, n_heads, window_size=4096):
        super().__init__()
        self.window_size = window_size
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, kv_cache=None):
        B, T, _ = x.shape
        q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        if kv_cache is not None:
            k = torch.cat([kv_cache[0], k], dim=2)
            v = torch.cat([kv_cache[1], v], dim=2)

        # 按窗口大小限制 KV Cache
        if k.shape[2] > self.window_size:
            k = k[:, :, -self.window_size:]
            v = v[:, :, -self.window_size:]

        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_o(out), (k, v)

PagedAttention（vLLM）

基于虚拟内存的 KV Cache 管理

vLLM 的 PagedAttention 受操作系统虚拟内存系统的启发，以 固定大小的块（页） 为单位管理 KV Cache。由此消除了连续内存分配带来的低效（内部碎片、外部碎片）。

class PagedKVCache:
    """vLLM PagedAttention 概念实现
    以页为单位管理 KV Cache"""
    def __init__(self, block_size=16, n_blocks=1024, n_heads=32, d_k=128):
        self.block_size = block_size
        self.n_blocks = n_blocks
        self.n_heads = n_heads
        self.d_k = d_k

        # 物理块池（预先分配）
        self.k_blocks = torch.zeros(n_blocks, n_heads, block_size, d_k)
        self.v_blocks = torch.zeros(n_blocks, n_heads, block_size, d_k)

        # 空闲块管理
        self.free_blocks = list(range(n_blocks))

        # 按序列的块表（虚拟 -> 物理映射）
        self.block_tables = {}  # seq_id -> list of physical block indices
        self.seq_lengths = {}

    def allocate(self, seq_id):
        """为新序列分配首个块"""
        if not self.free_blocks:
            raise RuntimeError("OOM: No free blocks available")
        block_idx = self.free_blocks.pop(0)
        self.block_tables[seq_id] = [block_idx]
        self.seq_lengths[seq_id] = 0

    def append_token(self, seq_id, k, v):
        """向序列追加新 token 的 KV"""
        seq_len = self.seq_lengths[seq_id]
        block_idx_in_seq = seq_len // self.block_size
        offset = seq_len % self.block_size

        # 当前块已满时分配新块
        if offset == 0 and block_idx_in_seq > 0:
            if not self.free_blocks:
                raise RuntimeError("OOM: No free blocks available")
            new_block = self.free_blocks.pop(0)
            self.block_tables[seq_id].append(new_block)

        # 把 KV 写入物理块
        physical_block = self.block_tables[seq_id][block_idx_in_seq]
        self.k_blocks[physical_block, :, offset] = k
        self.v_blocks[physical_block, :, offset] = v
        self.seq_lengths[seq_id] += 1

    def get_kv(self, seq_id):
        """拼出该序列完整的 KV Cache"""
        blocks = self.block_tables[seq_id]
        seq_len = self.seq_lengths[seq_id]

        k_list, v_list = [], []
        for i, block_idx in enumerate(blocks):
            if i == len(blocks) - 1:
                # 最后一个块只取实际 token 数
                remaining = seq_len % self.block_size or self.block_size
                k_list.append(self.k_blocks[block_idx, :, :remaining])
                v_list.append(self.v_blocks[block_idx, :, :remaining])
            else:
                k_list.append(self.k_blocks[block_idx])
                v_list.append(self.v_blocks[block_idx])

        return torch.cat(k_list, dim=1), torch.cat(v_list, dim=1)

    def free(self, seq_id):
        """序列结束后归还块"""
        for block_idx in self.block_tables[seq_id]:
            self.free_blocks.append(block_idx)
        del self.block_tables[seq_id]
        del self.seq_lengths[seq_id]

PagedAttention 的主要优点如下。

消除内存碎片：不需要连续内存，因此不存在外部碎片
提升内存使用效率：只按实际用量分配块。vLLM 相比传统方式最高可实现 24 倍的吞吐量
Copy-on-Write：在 beam search 或并行采样中共享 KV Cache，从而节省内存
前缀缓存（Prefix Caching）：多个请求之间共享公共前缀（系统提示词）的 KV Cache

故障案例与恢复流程

案例 1：长上下文推理中发生 OOM

情况：使用 Llama 3 70B 模型尝试对 128K 上下文长度的文档做摘要，在批大小为 4 时发生 OOM，导致整个推理服务中断。

内存分析：

模型权重（FP16）：约 140 GB
每个请求的 KV Cache（GQA, H=64, G=8）：2 x 80 x 8 x 128K x 128 x 2 = 约 33.5 GB
批大小 4 的 KV Cache 总量：约 134 GB
所需总内存：约 274 GB（即便是 8x A100 80GB = 640 GB，考虑到激活内存与临时缓冲区，也存在风险）

恢复流程：

# 1. 重启 vLLM 服务器时设置内存限制
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-70B \
  --tensor-parallel-size 8 \
  --max-model-len 65536 \
  --gpu-memory-utilization 0.9 \
  --max-num-seqs 8 \
  --enforce-eager

# 2. 启用 KV Cache 量化
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-70B \
  --tensor-parallel-size 8 \
  --kv-cache-dtype fp8_e5m2 \
  --max-model-len 128000

# 3. 动态限制批大小
# 减小 max-num-seqs 以限制并发请求数

案例 2：KV Cache 量化导致质量下降

情况：为降低推理成本把 KV Cache 量化为 INT4，但在长对话中响应质量急剧下降。尤其是在数学计算和代码生成任务中，错误率明显上升。

症状：

对话超过 10 轮后 hallucination 增多
数学题准确率下降 15%
代码生成时频繁出现语法错误

恢复流程：

# 1. 把量化精度从 INT4 提升到 INT8
# 在 vLLM 中修改 kv-cache-dtype
# --kv-cache-dtype fp8_e5m2  (使用 FP8)

# 2. 只在关键层维持高精度（混合量化）
# 首尾层用 FP16，中间层用 INT8
kv_cache_config = {
    "default_dtype": "int8",
    "high_precision_layers": [0, 1, 2, 77, 78, 79],  # 首/尾各 3 层
    "high_precision_dtype": "fp16"
}

# 3. 重新跑基准测试以确认质量
# 用 MMLU、HumanEval、GSM8K 等标准基准做验证

案例 3：Prefix Caching 出错导致响应串扰

情况：启用 vLLM 的前缀缓存功能后，出现了不同用户请求被套用错误上下文的问题。系统提示词相同但会话不同的请求之间，KV Cache 被错误地共享了。

恢复流程：

# 1. 暂时禁用前缀缓存
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-70B \
  --enable-prefix-caching false

# 2. 修改前缀哈希键，使其包含会话 ID
# 即便前缀相同，不同会话也各自使用独立的缓存

# 3. 重新启用前缀缓存时设置 TTL
# --prefix-cache-ttl 300  (5 分钟后自动过期)

优化清单

模型选择与设置

确认模型采用的注意力机制（MHA/MQA/GQA/MLA）
使用 GQA 模型时验证 n_kv_heads 的最优值（通常为 n_heads/8）
按实际使用模式限制最大序列长度

KV Cache 内存管理

计算 KV Cache 占 GPU 内存的分配比例（建议为总量的 60-80%）
决定是否应用 KV Cache 量化（品质/内存权衡：FP8 > INT8 > INT4）
决定是否引入 PagedAttention（vLLM、TensorRT-LLM 等）
启用前缀缓存时验证是否能防止键冲突

推理性能优化

启用 Continuous Batching（请求完成后立即回收槽位）
评估是否引入 Speculative Decoding（用草稿模型最大化 KV Cache 利用率）
评估是否适用 Sliding Window Attention（长文档摘要等场景）
在 Tensor Parallelism 与 Pipeline Parallelism 之间做选择

监控

收集 KV Cache 使用率指标（vLLM: vllm:cache_usage_percent）
跟踪各请求的 KV Cache 内存占用
设置 OOM 事件自动告警
按批大小做吞吐量与延迟的性能剖析

运维注意事项

设置并发请求数上限（max-num-seqs）
设置 GPU 内存使用率上限（gpu-memory-utilization）
实现 OOM 发生时的优雅降级（请求排队、缩小批大小）
定期做 KV Cache 剖析以确认是否存在内存泄漏

结语

KV Cache 优化是 LLM 推理效率的核心。从 MHA 到 MQA、GQA、MLA 的注意力机制演进，成功在维持模型质量的同时大幅削减了 KV Cache 内存。GQA 是在 Llama 系列中得到验证的实用方案，而 MLA 正如 DeepSeek-V2/V3 所展示的，让更激进的压缩成为可能。

KV Cache 量化、驱逐策略、Sliding Window Attention，是在内存受限条件下处理长序列的补充技术；PagedAttention 则从内存管理本身入手革新，最大化了推理系统的吞吐量。

在实际生产环境中，这些技术需要组合使用，而要找到匹配模型特性、工作负载模式与 GPU 资源的最优配置，持续的性能剖析与基准测试是不可或缺的。

引言

Transformer Self-Attention 与 KV Cache 基础

Self-Attention 运算

KV Cache 内存分析

Multi-Head Attention（MHA）内存分析

MHA 结构

Multi-Query Attention（MQA）

MQA 结构与节省效果

Grouped Query Attention（GQA）

GQA 架构（Llama 2/3）

Multi-Head Latent Attention（MLA）

MLA 架构（DeepSeek-V2/V3）

MHA vs MQA vs GQA vs MLA 对比表

KV Cache 压缩技术

量化（Quantization）

驱逐策略（Eviction Policy）

Sliding Window Attention

PagedAttention（vLLM）

基于虚拟内存的 KV Cache 管理

故障案例与恢复流程

案例 1：长上下文推理中发生 OOM

案例 2：KV Cache 量化导致质量下降

案例 3：Prefix Caching 出错导致响应串扰

优化清单

模型选择与设置

KV Cache 内存管理

推理性能优化

监控

运维注意事项

结语

参考资料