大模型底层原理：注意力机制优化与长上下文处理

一、注意力机制的计算瓶颈与长上下文的工程挑战

Transformer 架构的核心——自注意力机制（Self-Attention）的计算复杂度为 O(n²)，其中 n 为序列长度。这意味着当上下文窗口从 4K 扩展到 128K 时，注意力计算量增长约 1000 倍。在实际推理中，一个 128K 上下文的请求可能消耗 40GB 以上的显存，推理延迟从毫秒级飙升到分钟级。

这种计算瓶颈直接限制了 AI 产品的商业化落地。在 RAG 场景中，检索到的文档片段可能达到数万 Token；在代码辅助场景中，项目上下文可能超过 10 万 Token。如果模型无法高效处理长上下文，这些场景就只能依赖截断或摘要，导致信息丢失和输出质量下降。

二、注意力机制的数学原理与优化路径

2.1 标准自注意力的计算流程

标准自注意力的计算分为三步：线性投影生成 Q/K/V、注意力分数计算、加权求和。

graph LR
    A[输入 X] --> B[线性投影: Q=XWq, K=XWk, V=XWv]
    B --> C[注意力分数: S=QK^T / √d]
    C --> D[Softmax 归一化: A=softmax S]
    D --> E[加权求和: Output=AV]
    F[KV Cache] --> C
    G[位置编码] --> B

其中 QK^T 的计算是瓶颈所在。对于序列长度 n 和头维度 d，QK^T 产生一个 n×n 的注意力矩阵，需要 O(n²d) 的计算量和 O(n²) 的存储空间。

2.2 四种主流优化策略

KV Cache：在自回归推理中，已生成的 Token 的 K/V 不需要重复计算，只需缓存并在后续步骤中复用。这是最基础也最有效的优化，将推理复杂度从 O(n²) 降低到 O(n)（单步推理）。但 KV Cache 本身占用大量显存——一个 7B 模型在 128K 上下文下，KV Cache 可能占用 16GB 显存。

Flash Attention：通过分块计算（Tiling）和内核融合（Kernel Fusion），避免在 GPU HBM 中实例化完整的 n×n 注意力矩阵。Flash Attention 将注意力计算拆分为适合 SRAM 的小块，逐块计算后累加结果，显存占用从 O(n²) 降低到 O(n)。这是目前最广泛采用的优化方案。

MQA/GQA：Multi-Query Attention 让所有注意力头共享同一组 K/V 投影，仅 Q 保持多头。Grouped-Query Attention 是 MQA 的折中方案，将多个头归为一组共享 K/V。GQA 在几乎不损失模型质量的前提下，将 KV Cache 大小减少到原来的 1/8~1/4。

稀疏注意力：只计算部分 Token 对之间的注意力分数，跳过不重要的连接。典型方案包括滑动窗口注意力（仅关注邻近 Token）和全局注意力（少量关键 Token 与所有 Token 计算注意力）。稀疏注意力将计算复杂度降低到 O(n×w)，其中 w 为窗口大小。

三、长上下文处理的工程实现

3.1 KV Cache 管理与显存优化

from dataclasses import dataclass
from typing import Optional
import math


@dataclass
class KVCacheConfig:
    """KV Cache 配置参数"""
    num_layers: int          # 模型层数
    num_heads: int           # 注意力头数
    head_dim: int            # 每个头的维度
    num_kv_heads: int        # KV 头数（GQA 时小于 num_heads）
    max_seq_len: int         # 最大序列长度
    dtype_bytes: int = 2     # FP16 每个参数占 2 字节

    @property
    def cache_size_per_token(self) -> int:
        """每个 Token 的 KV Cache 大小（字节）"""
        # 每层: 2(K+V) × num_kv_heads × head_dim
        return 2 * self.num_kv_heads * self.head_dim * self.num_layers * self.dtype_bytes

    @property
    def max_cache_size(self) -> int:
        """最大序列长度下的 KV Cache 总大小"""
        return self.cache_size_per_token * self.max_seq_len

    def estimate_gpu_memory(self, model_params_gb: float) -> dict:
        """估算推理所需 GPU 显存"""
        cache_gb = self.max_cache_size / (1024 ** 3)
        total = model_params_gb + cache_gb
        return {
            "model_params_gb": model_params_gb,
            "kv_cache_gb": round(cache_gb, 2),
            "total_gb": round(total, 2),
            "recommendation": self._gpu_recommendation(total),
        }

    def _gpu_recommendation(self, total_gb: float) -> str:
        if total_gb <= 24:
            return "单卡 A10G (24GB) 或 RTX 4090 (24GB)"
        elif total_gb <= 48:
            return "单卡 A6000 (48GB) 或 2×A10G"
        elif total_gb <= 80:
            return "单卡 A100 (80GB)"
        else:
            return "多卡 A100 或使用量化降低显存"


# 示例：Qwen2.5-7B 的 KV Cache 估算
config = KVCacheConfig(
    num_layers=28,
    num_heads=28,
    head_dim=128,
    num_kv_heads=4,  # GQA: 4 组 KV 头
    max_seq_len=131072,  # 128K 上下文
)

# 估算结果
estimate = config.estimate_gpu_memory(model_params_gb=14.0)
# KV Cache 约 7.0GB，总显存约 21GB → 单卡 24GB 可运行

3.2 滑动窗口注意力实现

import torch
import torch.nn.functional as F


class SlidingWindowAttention(torch.nn.Module):
    """滑动窗口注意力，仅计算窗口内的 Token 对"""

    def __init__(self, dim: int, num_heads: int, window_size: int = 256):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.window_size = window_size

        self.q_proj = torch.nn.Linear(dim, dim, bias=False)
        self.k_proj = torch.nn.Linear(dim, dim, bias=False)
        self.v_proj = torch.nn.Linear(dim, dim, bias=False)
        self.out_proj = torch.nn.Linear(dim, dim, bias=False)

    def forward(self, x: torch.Tensor, kv_cache: Optional[tuple] = None):
        batch_size, seq_len, _ = x.shape

        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
        k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
        v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)

        # 拼接 KV Cache（自回归推理时）
        if kv_cache is not None:
            past_k, past_v = kv_cache
            k = torch.cat([past_k, k], dim=1)
            v = torch.cat([past_v, v], dim=1)

        # 构建滑动窗口掩码
        total_len = k.shape[1]
        mask = torch.ones(seq_len, total_len, dtype=torch.bool, device=x.device)
        for i in range(seq_len):
            # 当前 Token 可以关注窗口范围内的历史 Token
            query_pos = i + (total_len - seq_len)  # 绝对位置
            window_start = max(0, query_pos - self.window_size + 1)
            mask[i, :window_start] = False  # 窗口外的位置被屏蔽

        # 转置为 [batch, heads, seq, dim] 以适配 scaled_dot_product_attention
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # 使用 PyTorch 2.0+ 的 Flash Attention 实现
        output = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=mask.unsqueeze(0).unsqueeze(0).expand(
                batch_size, self.num_heads, -1, -1
            ),
            is_causal=False,
        )

        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        return self.out_proj(output), (k.transpose(1, 2), v.transpose(1, 2))

四、注意力优化的工程权衡

精度与效率的取舍：MQA/GQA 通过减少 KV 头数降低显存和计算量，但可能影响模型在复杂推理任务上的表现。实测数据显示，GQA 在大多数基准测试上与 MHA 差距在 1%~2% 以内，但在需要精细注意力分布的任务（如长文档问答）上差距可能扩大到 3%~5%。选择 GQA 组数时需要在显存预算和精度要求之间找到平衡。

稀疏注意力的信息损失：滑动窗口注意力假设远距离 Token 的依赖关系较弱，但这一假设在某些场景下不成立——例如法律文档中，定义条款可能出现在文档开头，而引用出现在末尾。纯滑动窗口方案会丢失这类长距离依赖。Mistral 的解决方案是滚动缓冲区（Rolling Buffer），配合少量全局注意力 Token 来捕获关键信息。

KV Cache 的显存竞争：在多用户并发推理时，不同请求的 KV Cache 共享 GPU 显存。当显存不足时，需要驱逐某些请求的 Cache，导致下次推理需要重新计算。PagedAttention（vLLM 的核心创新）通过虚拟内存管理解决了这一问题，将 KV Cache 分页存储，按需分配和回收。

Flash Attention 的硬件依赖：Flash Attention 需要 GPU 的 SRAM 容量足够大来容纳分块计算的数据。不同 GPU 架构的 SRAM 大小不同，A100 的 SRAM 为 192MB，而 V100 仅为 32MB。在 SRAM 不足的 GPU 上，Flash Attention 需要更小的分块尺寸，性能优势会打折扣。

五、总结

注意力机制的优化是长上下文处理的核心工程挑战。KV Cache 是推理加速的基础，Flash Attention 解决了显存瓶颈，GQA 在精度与效率间取得平衡，稀疏注意力为超长序列提供了可行方案。在工程落地时，需要根据 GPU 显存预算、上下文长度需求和精度要求选择合适的优化组合：4K 上下文用标准 MHA + KV Cache 即可，32K 上下文推荐 GQA + Flash Attention，128K 以上需要叠加稀疏注意力和 PagedAttention。关键原则是：优化不是免费的，每种优化都伴随着精度或灵活性的代价，需要通过基准测试验证在目标场景下的实际效果。