4、MHA性能调优（下）：Sparse Attention、Linear Attention与KV Cache优化策略

上一讲我们聊了MHA的计算瓶颈和Flash Attention。今天继续往下挖，把剩下的三板斧讲透：Sparse Attention、Linear Attention，还有KV Cache的优化。这三个方向，说白了都是在跟“注意力矩阵”的平方复杂度死磕。

我个人习惯把优化策略分成两类：一类是“偷懒”，比如稀疏化，只算重要的部分；另一类是“变形”，比如线性注意力，把计算顺序彻底换掉。KV Cache优化则更偏工程，属于“存得好不如算得巧”。

4.1 Sparse Attention：别算那些没用的

你想想看，一个长序列里，每个token真的需要跟所有其他token都交互吗？我在项目中遇到过，处理一篇5000字的文档，大部分token之间的注意力分数都趋近于零。算它们纯粹是浪费算力。

Sparse Attention的核心思路就是：只保留最重要的连接。常见的模式有几种：

  • 局部窗口注意力：每个token只看前后k个邻居。适合局部语义强的场景，比如代码、短文本。
  • 全局稀疏注意力：选一些“锚点”token（比如[CLS]），让它们跟所有token交互，其他token只跟局部邻居交互。BigBird和Longformer就是典型代表。
  • 随机稀疏注意力：随机采样一部分位置做交互。理论上能近似全连接的效果，但实际工程中不太常用，因为随机访问对GPU不友好。

关键点：稀疏模式的选择取决于你的数据特性。文本生成任务，局部窗口就够用；长文档分类，全局稀疏更合适。

代码实现上，稀疏注意力通常需要自定义CUDA kernel。不过现在主流框架都提供了封装，比如PyTorch的torch.nn.TransformerEncoderLayer支持传入src_mask来实现稀疏化。但说实话，直接用mask做稀疏，效率并不高——因为底层还是算了完整的注意力矩阵，只是把结果mask掉了。

# 伪代码：稀疏注意力掩码示例
def create_sparse_mask(seq_len, window_size):
    mask = torch.zeros(seq_len, seq_len)
    for i in range(seq_len):
        start = max(0, i - window_size)
        end = min(seq_len, i + window_size + 1)
        mask[i, start:end] = 1
    return mask.bool()

注意：这种mask方式只适合小规模实验。真正工业级部署，必须用Block Sparse Attention或者Triton手写kernel，否则内存和计算量一点没省。

4.2 Linear Attention：把平方复杂度降成线性

Sparse Attention再怎么优化，本质上还是O(n²)的复杂度。Linear Attention则彻底换了个思路——用核技巧把注意力计算拆解成矩阵乘法的结合律。

标准注意力公式是：Attention(Q,K,V) = softmax(QK^T)V。这里的softmax让计算无法交换顺序。Linear Attention的做法是：去掉softmax，或者用其他核函数近似，让计算变成 (QK^T)V = Q(K^TV)。

你想想看，K^TV这一步是固定的，可以预先算好。这样复杂度就从O(n²d)降到了O(nd²)。当序列长度n远大于特征维度d时，效果非常明显。

我在项目中用过Performer（用随机特征映射近似softmax），也试过Linear Transformer（直接用ELU激活函数替代softmax）。说实话，Linear Attention在长序列任务上确实快，但精度往往有损失。尤其是需要精细对齐的任务，比如机器翻译，效果不如标准MHA。

方法 复杂度 精度损失 适用场景
标准MHA O(n²d) 无 短序列、高精度任务
Sparse Attention O(nk) (k为窗口大小) 较小 长文本、局部依赖强
Linear Attention O(nd²) 中等 超长序列、实时推理

我的建议：如果序列长度在1024以内，别折腾Linear Attention，直接用Flash Attention。超过2048，再考虑稀疏或线性方案。

4.3 KV Cache优化：推理阶段的命门

终于聊到KV Cache了。这是自回归生成任务里最关键的优化点。说白了，KV Cache就是存下之前所有token的Key和Value矩阵，避免重复计算。

举个例子：生成第100个token时，标准做法是重新算一遍前99个token的注意力。有了KV Cache，只需要算第100个token的Query，然后跟缓存的K、V做注意力就行。这能省掉大约90%的计算量。

但KV Cache也有代价——显存占用。假设模型有32层，每层hidden size是4096，序列长度2048，batch size是1，那么KV Cache的大小大约是：32 × 2 × 2048 × 4096 × 2字节 ≈ 1GB。如果batch size变成64，就是64GB。显存直接爆炸。

我曾经在一个对话系统项目里，就因为KV Cache没优化好，导致推理时显存溢出。后来用了几个技巧才搞定：

  • 共享KV Cache：多轮对话中，历史轮次的KV Cache可以复用，不需要重新算。
  • 量化KV Cache：用INT8甚至INT4存储K和V，精度损失很小，但显存能省一半以上。
  • 分页KV Cache：像操作系统管理内存一样，把KV Cache分页存储，按需加载。vLLM就是靠这个技术出名的。
# KV Cache的典型实现（伪代码）
class KVCache:
    def __init__(self, max_batch, max_seq, num_layers, head_dim):
        self.cache = [None] * num_layers  # 每层一个缓存
    
    def update(self, layer_idx, key, value):
        if self.cache[layer_idx] is None:
            self.cache[layer_idx] = (key, value)
        else:
            old_k, old_v = self.cache[layer_idx]
            self.cache[layer_idx] = (
                torch.cat([old_k, key], dim=2),
                torch.cat([old_v, value], dim=2)
            )
    
    def get(self, layer_idx):
        return self.cache[layer_idx]

避坑指南：我曾经犯过一个错误——在生成每个token时都重新初始化KV Cache。结果显存占用翻倍，速度反而更慢。正确的做法是：在生成开始时初始化一次，后续只做append操作。

4.4 三种策略的联合使用

实际项目中，这三种策略往往不是互斥的。我见过一个工业级方案：

  1. 用Sparse Attention减少单次计算量（窗口大小设为512）
  2. 用Linear Attention处理窗口外的长距离依赖（用Performer近似）
  3. 用KV Cache量化降低显存占用（INT8量化，精度损失<0.5%）

这个方案在128K长度的文档摘要任务上，推理速度提升了4倍，显存占用降低了60%。当然，精度也掉了1%左右。但业务方觉得可以接受。

嗯，这里要注意：不要为了优化而优化。如果你的序列长度只有512，用标准MHA+Flash Attention就足够了。优化是有成本的——代码复杂度、调试难度、精度损失，这些都要权衡。

最后总结一下我的个人经验：

  • 序列长度 < 1024：用Flash Attention，别折腾
  • 序列长度 1024~8192：考虑Sparse Attention + KV Cache量化
  • 序列长度 > 8192：Linear Attention + 分页KV Cache

下一讲我们会聊MHA的分布式训练优化，包括张量并行和序列并行。到时候再细聊怎么把注意力计算分布到多卡上。

← 上一章 📖 返回目录 下一章 →