Transformer高效注意力机制实现：Flash Attention与GQA原理对比与代码实践

标准多头注意力机制的实现与计算瓶颈

Transformer架构中的多头注意力（Multi-Head Attention, MHA）是自然语言处理模型的核心组件。标准MHA的每个注意力头拥有独立的Query、Key、Value投影矩阵，计算复杂度为O(n²d)，其中n为序列长度，d为模型维度。当序列长度超过8K时，显存占用和计算延迟急剧上升，成为大模型推理和训练的主要瓶颈。

PyTorch实现标准MHA的核心代码如下：

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_model = d_model
        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):
        batch, seq_len, _ = x.shape
        Q = self.w_q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.w_k(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.w_v(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)
        return self.w_o(out)

在MHA中，KV缓存的显存占用与注意力头数成正比。以Llama 2 70B为例，32个注意力头在batch_size=32、seq_len=4096时，KV缓存需占用约40GB显存，严重制约推理吞吐量。减少KV缓存大小成为注意力机制优化的核心方向。

Grouped-Query Attention(GQA)原理与实现

GQA通过共享Key和Value投影矩阵来减少KV缓存大小。在MHA中，每个注意力头有独立的K和V；在GQA中，多个Query头共享一组K和V。当共享组数等于1时，退化为Multi-Query Attention（MQA）；当共享组数等于头数时，即为标准MHA。GQA在两者之间取得平衡。

GQA的PyTorch实现：

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_rep = n_heads // n_kv_heads
        self.d_k = d_model // n_heads
        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 forward(self, x):
        batch, seq_len, _ = x.shape
        Q = self.w_q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.w_k(x).view(batch, seq_len, self.n_kv_heads, self.d_k).transpose(1, 2)
        V = self.w_v(x).view(batch, seq_len, self.n_kv_heads, self.d_k).transpose(1, 2)
        # 扩展KV头以匹配Q头数量
        K = K.repeat_interleave(self.n_rep, dim=1)
        V = V.repeat_interleave(self.n_rep, dim=1)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
        return self.w_o(out)

GQA将KV缓存减少为MHA的n_kv_heads/n_heads倍。Llama 2 70B使用8个KV头（32个Q头的1/4），KV缓存显存降低75%，推理速度提升约1.7倍，质量损失可忽略。GQA已作为默认配置应用于Llama 3、Mistral、DeepSeek V3等主流开源模型。

Flash Attention IO优化原理与代码实践

Flash Attention通过分块计算（tiling）和内核融合（kernel fusion）优化注意力计算的标准HBM读写次数。标准注意力的中间矩阵S（n×n）和P（n×n）需要反复在HBM和SRAM之间搬运。Flash Attention将Q、K、V分块加载到SRAM中，在SRAM内完成注意力计算后直接写回结果，避免中间矩阵的全量HBM读写。

使用Flash Attention库的实现：

from flash_attn import flash_attn_func

class FlashAttentionLayer(nn.Module):
    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)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch, seq_len, _ = x.shape
        Q = self.w_q(x).view(batch, seq_len, self.n_heads, self.d_k)
        K = self.w_k(x).view(batch, seq_len, self.n_heads, self.d_k)
        V = self.w_v(x).view(batch, seq_len, self.n_heads, self.d_k)
        # flash_attn_func 接收 [batch, seq, heads, dim] 格式
        out = flash_attn_func(Q, K, V, causal=True)
        out = out.view(batch, seq_len, -1)
        return self.w_o(out)

Flash Attention在序列长度16K时，相比标准实现显存占用从64GB降至2GB，训练速度提升2-4倍。该算法的IO复杂度从O(n²)降低到O(n²M⁻¹)，M为SRAM大小，且计算结果与标准注意力数学等价，无精度损失。Flash Attention 2进一步优化了并行度分配，将序列维度并行化，在长序列场景下GPU利用率提升显著。

注意力机制变体对比与选型建议

三种注意力机制的核心差异在于KV头数量和IO优化策略。MHA每个Q头独立对应KV头，精度最高但显存占用最大。GQA通过共享KV头减少缓存，精度接近MHA，适合推理场景。Flash Attention优化HBM读写但不改变KV头数量，可与GQA叠加使用。

选型建议：训练阶段优先使用Flash Attention + MHA组合，保证精度同时加速训练。推理阶段使用GQA + Flash Attention，显著降低KV缓存显存并提升吞吐量。对于边缘部署场景，MQA（GQA的极端形式）可将KV缓存降至最低，适合资源受限设备。实际工程中，vLLM和SGLang等推理框架已内置GQA和Flash Attention支持，通过配置参数即可启用。DeepSeek V3采用的MLA（Multi-head Latent Attention）进一步将KV缓存压缩到低维潜在空间，在128K上下文窗口下KV缓存仅占原始MHA的5%，是目前最激进的缓存优化方案。