注意力机制全家桶：从 Multi-Head 到 GQA 再到 Flash Attention 的手写实现

一、引言：注意力的进化史

2017年，Google 在 "Attention Is All You Need" 中提出了 Transformer 架构，其中 Multi-Head Attention（多头注意力）是核心组件。此后七年时间里，注意力机制经历了多次重要的演进：从最初的 Scaled Dot-Product Attention 到 Multi-Head Attention（MHA），再到 Multi-Query Attention（MQA）、Grouped Query Attention（GQA）、Flash Attention（闪速注意力）和各种线性近似注意力。每一次改进都不是锦上添花，而是在直面注意力机制固有的计算瓶颈。

注意力机制面临的核心矛盾是：表达能力强（全局交互）与计算代价高（O(N²) 复杂度）之间的根本矛盾。当序列长度 N 从 128 增长到 128K（如 DeepSeek 等前沿模型支持的超长上下文），注意力矩阵的规模从 16K 膨胀到 16G 个元素------这还不考虑 batch size 和头数。这个令人望而生畏的复杂度，正是 GQA 和 Flash Attention 这些技术所要解决的问题。

理解这些注意力变体，对于任何需要处理大规模语言模型的工程师来说都至关重要。不是你选择要不要用的问题，而是你的模型部署和使用是否高效的问题。如果你在部署 LLaMA 系列、Falcon 系列或 DeepSeek 系列模型，这些架构细节直接决定了你能否在单张 GPU 上运行长序列推理。

本文将从最基础的 Scaled Dot-Product Attention 开始，逐步深入到 MHA、GQA、MQA 的差异，最后手写简化版的 Flash Attention。每个实现都附带完整的 PyTorch 代码、复杂度分析和实际使用场景。

技术预备：Python 3.8+、PyTorch 2.x、理解 Transformer 基本概念。如果熟悉 Softmax 和矩阵乘法的基础知识就足够了。

---

二、基础回顾：Scaled Dot-Product Attention

2.1 数学原理

在讨论注意力机制的改进变体之前，有必要先牢固掌握基础版本的数学原理。Scaled Dot-Product Attention 是 Transformer 一切变体的起点，理解它的每一处设计细节，后续的 GQA 和 Flash Attention 才能水到渠成。

给定 Query（查询）、Key（键）、Value（值）三个矩阵，注意力的计算分为三步：

1. 计算相关性：Q 与 K 的转置做矩阵乘法，得到注意力分数矩阵，表示每个 query 与每个 key 的相关性。在这一步，如果 query 是当前 token，key 是所有历史 token，那么注意力分数就代表了"当前 token 应该关注哪些历史 token"。

2. 缩放与归一化：除以 sqrt(d_k) 防止梯度消失，然后做 Softmax 归一化，将分数转换为概率分布。这一步决定了每个 token 最终贡献多少信息到当前 token 的表示中。

3. 加权求和：用归一化后的注意力权重对 Value 做加权平均，得到最终的输出。注意力权重越高的 token，其 Value 对最终输出的贡献越大。

公式如下：Attention(Q,K,V) = softmax(QK^T / sqrt(d_k)) V

其中 Q ∈ ℝ^(N×d_k)，K ∈ ℝ^(M×d_k)，V ∈ ℝ^(M×d_v)。通常 N = M，即 query 和 key/value 序列长度相同。在自注意力（Self-Attention）中，Q、K、V 都来自同一个输入序列。

2.2 为什么需要缩放

除以 sqrt(d_k) 这一步看似简单，但背后有深刻的统计意义。假设 Q 和 K 中的每个元素都是从标准正态分布 N(0,1) 中独立采样的，那么 Q 和 K 的点积就是 d_k 个独立随机变量的和。根据概率论的中心极限定理，这个和的均值是 0，方差是 d_k。这意味着点积的标准差为 sqrt(d_k)，结果分布在 -3√d_k, 3√d_k 的范围内。

对于 d_k=512（标准 Transformer 中 head 的常用维度），这个范围大约是 -68, 68。Softmax 在这样的输入值下几乎退化为 one-hot 形式：最大值对应概率接近 1，其余值接近 0。这不仅导致信息几乎完全丢失（只有一个位置的注意力被

除以 sqrt(d_k) 后，方差恢复到 1，Softmax 的输入值分布在合理的 -3, 3 区间内，梯度就可以正常流动了。这个简单的缩放操作是训练深层 Transformer 的关键之一。

2.3 基础实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    基础 Scaled Dot-Product Attention

    参数:
        Q: (B, H, N, d_k) --- 查询
        K: (B, H, M, d_k) --- 键
        V: (B, H, M, d_v) --- 值
        mask: (B, 1, N, M) 或 (N, M) --- 注意力掩码

    返回:
        output: (B, H, N, d_v)
        attention_weights: (B, H, N, M) --- 可用于可视化
    """
    d_k = Q.size(-1)

    # 1. QK^T 矩阵乘法
    scores = torch.matmul(Q, K.transpose(-2, -1))  # (B, H, N, M)

    # 2. 缩放
    scores = scores / math.sqrt(d_k)

    # 3. 应用掩码（causal mask 防止看到未来 token）
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # 4. Softmax 归一化
    attention_weights = F.softmax(scores, dim=-1)

    # 5. 加权求和
    output = torch.matmul(attention_weights, V)

    return output, attention_weights

复杂度分析 ：

• 时间复杂度：O(N·M·d) = O(N²·d)（当 N=M 时）。每一步矩阵乘法都是 O(N²·d)。

• 空间复杂度：O(N²)------需要存储完整的注意力矩阵。对于 N=4096，就是 4096²=16M 个元素。

这就是注意力机制的根本瓶颈。当 N=4096 时，注意力矩阵有 16M 个元素；当 N=128K 时，这个数字达到 16G------仅仅存储注意力矩阵就需要 32GB 的显存（fp16）。这还是在单层、单头的情况。实际模型中 Layer × Head 会让这个数字再大几个数量级。例如一个 32 层、32 头的模型，N=8192 时仅注意力矩阵就需要 32×32×8192²×2≈137GB------远超当前任何单张 GPU 的容量。即使我们只计算一次（而不存储中间结果），QK^T 的计算量本身也是 O(N²·d) 的，当 N 很大时对计算资源的消耗同样巨大。因此，注意力优化的必要性不言而喻------它不仅关乎显存，也关乎计算时间和能耗。

---

三、Multi-Head Attention（多头注意力）

3.1 核心思想

多头注意力的想法很朴素：一个注意力头不够，那就用多个。这个直觉来源于一个观察：在自然语言中，不同的词之间往往需要同时建立多种关系------主谓关系、动宾关系、修饰关系等。单个注意力头只能产生一组注意力权重，无法同时捕捉这些不同类型的关系。而如果你使用多组独立的注意力参数，每组（每个头）就可以专门关注一种类型的关系。

从数学角度看，MHA 做的事情是在多个低维子空间中并行计算注意力，然后将结果拼接起来。这种

通过使用多个独立的注意力头（每个头有自己的 Q、K、V 投影参数），模型可以在不同的子空间中并行学习不同类型的注意力模式。最终将所有头的输出拼接，通过输出投影混合信息。这类似于 CNN 中多个卷积核的作用：每个卷积核检测不同的特征（边缘、纹理、形状），多个核并行工作形成丰富的特征表示。

具体做法：将 d_model 维度的 Q、K、V 拆分为 h 个 head，每个 head 的维度为 d_k = d_model / h。每个 head 独立计算注意力，然后将所有 head 的输出拼接后投影回 d_model 维。本质上，MHA 是"分组计算 + 信息融合"。

3.2 完整实现

class MultiHeadAttention(nn.Module):
    """
    标准多头注意力（Multi-Head Attention）

    参数:
        d_model (int): 模型维度
        n_heads (int): 注意力头数
        dropout (float): Dropout 概率
        bias (bool): 线性层是否使用偏置
    """
    def __init__(self, d_model=512, n_heads=8, dropout=0.1, bias=False):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.q_proj = nn.Linear(d_model, d_model, bias=bias)
        self.k_proj = nn.Linear(d_model, d_model, bias=bias)
        self.v_proj = nn.Linear(d_model, d_model, bias=bias)
        self.out_proj = nn.Linear(d_model, d_model, bias=bias)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        B = query.shape[0]

        # 1. 线性投影
        Q = self.q_proj(query)
        K = self.k_proj(key)
        V = self.v_proj(value)

        # 2. 分头: (B, N, d_model) → (B, n_heads, N, d_k)
        Q = Q.view(B, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(B, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(B, -1, self.n_heads, self.d_k).transpose(1, 2)

        # 3. 注意力计算（所有头并行）
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)

        # 4. 合并多头: (B, n_heads, N, d_k) → (B, N, d_model)
        output = attn_output.transpose(1, 2).contiguous()
        output = output.view(B, -1, self.d_model)

        # 5. 输出投影
        output = self.out_proj(output)
        output = self.dropout(output)
        return output

3.3 MHA 的局限

在推理阶段，MHA 面临严重的显存问题。以 LLaMA-7B（d_model=4096, n_heads=32, d_k=128）为例，对于 N=2048 的序列：

• 注意力矩阵：32×2048×2048×2(fp16)=256MB/层，32层共8GB

• KV Cache：2×32×128×2=16KB/token，2048个token约32MB/层，32层约1GB

合计约 9GB 仅用于注意力相关的存储。当序列增长到 8192 时，注意力矩阵增大 16 倍，KV Cache 增大 4 倍。MHA 的显存瓶颈直接催生了 MQA 和 GQA 的出现。

---

四、Multi-Query Attention（MQA）

4.1 核心思想

MQA 是 MHA 的一种极端简化版本：保持多个 Query 头，但 Key 和 Value 只有一个共享头。所有 Query 头共享同一套 K 和 V 投影。这个设计的灵感来自一个实际观察：在大多数自回归解码场景中，不同注意力头之间的差异主要来自 Query 的不同投影，而 Key 和 Value 在各头之间高度相似。既然 K/V 本质上很接近，为每个头分别存储多套 K/V 在推理时是巨大的浪费。

MQA 最早由 Shazeer 在 "Fast Transformer Decoding: One Write-Head is All You Need"（2019）中正式提出。实验表明，在相同计算预算下，MQA 的模型质量与 MHA 几乎持平，但推理速度提升约 30-50%。

4.2 实现

class MultiQueryAttention(nn.Module):
    """
    多查询注意力（Multi-Query Attention）
    --- 多个 Q 头，但 K 和 V 共享单个头
    """
    def __init__(self, d_model=512, n_heads=8, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, self.d_k, bias=False)
        self.v_proj = nn.Linear(d_model, self.d_k, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        B, N = query.shape[0], query.shape[1]

        Q = self.q_proj(query)
        Q = Q.view(B, -1, self.n_heads, self.d_k).transpose(1, 2)

        K = self.k_proj(key).unsqueeze(1)  # (B, 1, M, d_k)
        V = self.v_proj(value).unsqueeze(1)  # (B, 1, M, d_k)

        attn_output, _ = scaled_dot_product_attention(Q, K, V, mask)

        output = attn_output.transpose(1, 2).contiguous()
        output = output.view(B, -1, self.d_model)
        output = self.out_proj(output)
        output = self.dropout(output)
        return output

4.3 MQA 的优缺点

优点 ：

• KV Cache 减少到 MHA 的 1/h（对于 n_heads=32，减少 97%）。LLaMA-7B 的 KV Cache 从 ~1GB 降至 ~32MB。

• 推理阶段的显存占用大幅降低，同一个 GPU 可以塞进更大的 batch size，直接提升吞吐量。

• 解码速度提升显著------自回归解码是 memory-bound 的，每步减少数据搬运量就是每步减少延迟。

缺点 ：

• 单个 K/V 头的表达能力有限，在需要细粒度对齐的精细任务上质量略有下降。

• 但 PaLM、Falcon 等超大模型（100B+）的实践表明，模型规模越大，MQA 与 MHA 的差距越不明显。

---

五、Grouped Query Attention（GQA）

5.1 核心思想

GQA 是 MHA 和 MQA 的折中方案：将查询头分成若干组，每组内的 Q 头共享一对 K/V 头。组数 n_groups 是一个可配置的超参数：

• n_groups = n_heads → MHA（每个 Q 头独享 K/V）
• n_groups = 1 → MQA（所有 Q 头共享 K/V）
• n_groups = 2, 4, 8 → 中间状态

这种设计的哲学是：以极小的质量损失换取大幅的 KV Cache 降低。GQA 论文（Ainslie et al., 2023）表明，当 n_groups=8 时（即 n_heads=32, n_kv_heads=8），在多种 benchmark 上的质量下降几乎可以忽略不计，而 KV Cache 降低了 75%。

GQA 已经成为当前大模型的主流选择。LLaMA 2-70B 采用 GQA（n_heads=64, n_kv_heads=8），LLaMA 3-8B 也采用 GQA（n_heads=32, n_kv_heads=8）。

5.2 实现

class GroupedQueryAttention(nn.Module):
    """
    分组查询注意力（Grouped Query Attention）
    --- 将 n_heads 个 Q 头分成 n_groups 组，每组共享一对 K/V
    """
    def __init__(self, d_model=512, n_heads=8, n_groups=4, dropout=0.1):
        super().__init__()
        assert n_heads % n_groups == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_groups = n_groups
        self.n_kv_heads = n_groups
        self.heads_per_group = n_heads // n_groups
        self.d_k = d_model // n_heads

        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, self.d_k * n_groups, bias=False)
        self.v_proj = nn.Linear(d_model, self.d_k * n_groups, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        B, N = query.shape[0], query.shape[1]
        M = key.shape[1]

        Q = self.q_proj(query)
        Q = Q.view(B, -1, self.n_heads, self.d_k).transpose(1, 2)

        K = self.k_proj(key)
        K = K.view(B, -1, self.n_groups, self.d_k).transpose(1, 2)

        V = self.v_proj(value)
        V = V.view(B, -1, self.n_groups, self.d_k).transpose(1, 2)

        # 扩展 K/V 以匹配 Q 的头数
        K = K.unsqueeze(2).expand(-1, -1, self.heads_per_group, -1, -1)
        K = K.reshape(B, self.n_heads, M, self.d_k)

        V = V.unsqueeze(2).expand(-1, -1, self.heads_per_group, -1, -1)
        V = V.reshape(B, self.n_heads, M, self.d_k)

        attn_output, _ = scaled_dot_product_attention(Q, K, V, mask)

        output = attn_output.transpose(1, 2).contiguous()
        output = output.view(B, -1, self.d_model)
        output = self.out_proj(output)
        output = self.dropout(output)
        return output

5.3 GQA 的 KV Cache 节省

以 d_model=4096, n_heads=32, d_k=128, seq_len=8192 为例：

配置	n_groups	KV Cache / 层	节省比例
MHA	32	32×8192×128×2×2B=128MB	---
GQA-8	8	8×8192×128×2×2B=32MB	75%
GQA-4	4	4×8192×128×2×2B=16MB	87.5%
GQA-2	2	2×8192×128×2×2B=8MB	93.75%
MQA	1	1×8192×128×2×2B=4MB	96.9%

以 32 层计算，MHA 需要 4GB KV Cache，GQA-8 仅需 1GB------直接决定了是否能在单卡上运行长序列推理。

主流模型的实际选择：

模型	注意力类型	n_heads	n_kv_heads	策略
GPT-3 175B	MHA	96	96	大模型硬扛 MHA
LLaMA-1 7B	MHA	32	32	早期方案
LLaMA-2 7B	GQA	32	8	折中方案
LLaMA-3 8B	GQA	32	8	延续折中
DeepSeek-V2	MLA	128	1(低秩)	激进优化
PaLM	MQA	16	1	极端压缩
Falcon-40B	MQA	64	1	极端压缩

---

六、KV Cache：推理加速的核心

6.1 什么是 KV Cache？

在自回归解码中，模型逐个生成 token。生成第 t+1 个 token 时，需要基于前 t 个 token 的序列计算注意力。朴素的做法是每次重新计算所有历史 token 的 K 和 V------但这极其低效：前 t-1 个 token 的 K 和 V 已经算过，而且结果完全相同。每步都重新计算，就像你每读一页书都重新从前言开始读起。

KV Cache 的核心思想是：存储之前所有 token 的 K 和 V 在显存中，每次只计算新 token 的 K 和 V，然后将新的 K、V 追加到缓存中。这样，注意力计算就退化成了新 token 的 Q 与缓存中的全局 K 做矩阵乘法。

显存占用公式 ：

总 KV Cache = 2 × n_layers × n_kv_heads × seq_len × d_k × bytes_per_element

以 LLaMA-3-8B（n_layers=32, n_kv_heads=8, d_k=128）为例：

• seq_len=4096, fp16：2×32×8×4096×128×2=0.54GB

• seq_len=8192, fp16：2×32×8×8192×128×2=1.07GB

• seq_len=32768, fp16：2×32×8×32768×128×2=4.29GB

• seq_len=131072, fp16：2×32×8×131072×128×2=17.2GB

可以看到，即使使用了 GQA，超长序列下 KV Cache 对显存的消耗仍然可观。

6.2 带 KV Cache 的推理实现

class CachedAttention(nn.Module):
    """带 KV Cache 的 GQA 注意力（推理时使用）"""
    def __init__(self, d_model=512, n_heads=8, n_groups=4, dropout=0.1):
        super().__init__()
        self.gqa = GroupedQueryAttention(d_model, n_heads, n_groups, dropout)
        self.kv_cache = None

    def reset_cache(self):
        """开始新序列前重置"""
        self.kv_cache = None

    def forward_with_cache(self, x):
        """
        推理阶段的前向传播
        x: (B, 1, d_model) --- 当前生成的 token
        """
        if self.kv_cache is None:
            k = self.gqa.k_proj(x)
            v = self.gqa.v_proj(x)
            self.kv_cache = {'k': k, 'v': v}
        else:
            new_k = self.gqa.k_proj(x)
            new_v = self.gqa.v_proj(x)
            self.kv_cache['k'] = torch.cat([self.kv_cache['k'], new_k], dim=1)
            self.kv_cache['v'] = torch.cat([self.kv_cache['v'], new_v], dim=1)

        return self.gqa(x, self.kv_cache['k'], self.kv_cache['v'])

6.3 实际提升

实测数据（vLLM 对 LLaMA-7B 的基准测试）：

方式	延迟 (ms/token)	吞吐量 (tokens/s)
无 KV Cache	~35	~28
有 KV Cache	~8	~125

加速比约 4.4 倍。序列越长，加速效果越明显------因为每步计算量从 O(N²) 降为 O(N)。

---

七、Flash Attention：IO 感知的注意力算法

7.1 问题在哪里

前面所有注意力实现有一个共同的性能瓶颈：它们都忽略了 GPU 的内存层次结构。标准的注意力计算流程是：

1. 从 HBM（~2TB/s）读取 Q、K、V 到 SRAM（~19TB/s）
2. 计算 S = QK^T，将结果写回 HBM
3. 从 HBM 重新读取 S，计算 P = softmax(S)，写回 HBM
4. 从 HBM 重新读取 P 和 V，计算 O = PV，写回 HBM

注意力矩阵 S 和 P 在 HBM 和 SRAM 之间反复搬运，占用了大量 HBM 带宽。以 N=2048、d_k=128 为例，S 的大小是 2048×2048×2B=8MB。乘以 n_heads×n_layers 后，HBM 读写量达到 GB 级别------而 GPU 的 HBM 带宽是有限的共享资源。

7.2 GPU 内存层次

层级	容量	带宽	延迟
HBM (全局显存)	40-80 GB	~2 TB/s	~800 cycles
L2 Cache	40 MB	~8 TB/s	~200 cycles
SRAM (共享内存)	192 KB/SM	~19 TB/s	~30 cycles

SRAM 的带宽是 HBM 的约 10 倍，但容量极小（每个 SM 仅 192KB）。因此关键在于能否将注意力计算完全限制在 SRAM 内。

7.3 在线 Softmax 原理

标准 Softmax 需要一次性看到整行元素计算分母。Flash Attention 使用 Online Softmax（Safe Softmax 的变体），通过递归方式逐 block 更新统计量：

1. 处理 block 1：记录 m₁=max(x₁), s₁=Σexp(x₁-m₁)
2. 处理 block 2：更新全局最大值 m₂=max(m₁, max(x₂))
3. 校正 block 1：s₁' = s₁ × exp(m₁-m₂)
4. 合并 block 2：s₂ = s₁' + Σexp(x₂-m₂)
5. 校正输出：o = o₁ × exp(m₁-m₂) + Σexp(x₂-m₂)·v₂

递归进行，最终结果与全局 Softmax 完全一致。

7.4 简化版实现

def flash_attention_forward(Q, K, V, block_size=32):
    """
    简化版 Flash Attention（展示核心 Tiling + Online Softmax）
    """
    B, H, N, d_k = Q.shape
    d_v = V.shape[-1]
    O = torch.zeros(B, H, N, d_v, device=Q.device)

    for q_start in range(0, N, block_size):
        q_end = min(q_start + block_size, N)
        q_block = Q[:, :, q_start:q_end, :]

        o_block = torch.zeros(B, H, q_end - q_start, d_v, device=Q.device)
        m_prev = torch.full((B, H, q_end - q_start, 1), float('-inf'), device=Q.device)
        l_prev = torch.zeros(B, H, q_end - q_start, 1, device=Q.device)

        for kv_start in range(0, N, block_size):
            kv_end = min(kv_start + block_size, N)
            k_block = K[:, :, kv_start:kv_end, :]
            v_block = V[:, :, kv_start:kv_end, :]

            s = torch.matmul(q_block, k_block.transpose(-2, -1)) / math.sqrt(d_k)
            m_local = s.max(dim=-1, keepdim=True).values
            m_new = torch.max(m_prev, m_local)

            p = torch.exp(s - m_new)
            l_local = p.sum(dim=-1, keepdim=True)

            correction = torch.exp(m_prev - m_new)
            l_new = correction * l_prev + l_local
            o_block = (o_block * correction) + torch.matmul(p, v_block)

            m_prev, l_prev = m_new, l_new

        O[:, :, q_start:q_end, :] = o_block / l_prev

    return O

7.5 实际效果

来自 FlashAttention 论文（A100, N=4096, d=64, fp16）：

实现	前向+反向	加速比
PyTorch 标准实现	60.7 ms	1x
Flash Attention	9.1 ms	6.7x
Flash Attention-2	5.2 ms	11.7x

7.6 PyTorch 2.0 内置支持

从 PyTorch 2.0 开始，Flash Attention 已内置在 F.scaled_dot_product_attention 中：

output = F.scaled_dot_product_attention(
    Q, K, V, attn_mask=mask, dropout_p=0.0, is_causal=True
)

硬件要求：SM80+（A100/H100/RTX 3090+），d_k 需为 16 或 32 的倍数。PyTorch 会自动选择最优实现（Flash Attention / Memory Efficient Attention / 原生）。

---

八、注意力机制的工程权衡

8.1 计算 vs 通信

理解注意力的性能瓶颈，首先要明白 GPU 上的一个基本事实：大多数非训练场景下，注意力的瓶颈在显存带宽（memory bandwidth）而非计算（compute）。这个结论可能违反直觉------我们常说深度学习是计算密集型的，为什么注意力反而受限于内存带宽？

原因在于注意力计算的本质是一系列矩阵乘法（GEMM），而 GPU 在 GEMM 上的计算速度远快于数据从 HBM 搬运到 SRAM 的速度。以 N=2048、d_k=128 为例：QK^T 的计算量约 1G FLOPs（2048×2048×128×2 乘加运算），而 A100 的峰值算力是 312 TFLOPS（稀疏模式），理论计算时间约 3.2 微秒。但数据搬运所需的时间却是：Q 和 K 约 4MB（2×2048×128×4B），在 HBM 带宽约 2TB/s 下需约 2 微秒，再加上注意力矩阵写回 HBM 约 32MB（2048×2048×2×4B），约 16 微秒。约 80% 的时间花在了数据搬运上而非计算上。

这彻底解释了为什么 Flash Attention 能有 6-11x 的加速：它减少了 HBM 的读写次数，而不是加速了矩阵乘法本身。当注意力矩阵不需要写回 HBM 时，那 16 微秒的 HBM 写时间就被省去了。同样的逻辑也解释了为什么 KV Cache 对推理如此重要：它避免了重复计算 K/V，减少了每步的 HBM 搬运量。

进一步地，这个分析揭示了另一个重要洞见：注意力在长序列场景下是 memory-bound 的，而在短序列场景下是 compute-bound 的 。交叉点通常在 N=512 左右（取决于 GPU 型号）。这意味着：

• 短序列（N<512）：优化计算的收益更大（如使用 Tensor Core）

• 长序列（N>2048）：优化数据搬运的收益更大（如 Flash Attention、GQA）

8.2 粗略基准估计

综合 A100 公开基准数据（fp16, B=1, 单层）：

序列长度	MHA (标准)	MHA (FA)	GQA-4 + FA	主要瓶颈
512	~0.3 ms	~0.2 ms	~0.15 ms	kernel launch 开销
1024	~1.2 ms	~0.4 ms	~0.35 ms	数据传输
2048	~5 ms	~1.0 ms	~0.8 ms	HBM带宽 (O(N²))
4096	~20 ms	~3.0 ms	~2.0 ms	HBM带宽
8192	~80 ms	~10 ms	~5.0 ms	HBM + 计算混合

核心趋势：序列长度翻倍，标准 MHA 的延迟增加约 4 倍（O(N²) 的体现），而 Flash Attention 和 GQA 的增长更平缓。在 N=8192 时，GQA+FA 比标准 MHA 快约 16 倍------这不是理论增益，是实打实的推理加速。

8.3 组合使用的最佳实践

在实际的大模型推理系统中，GQA、KV Cache 和 Flash Attention 不是三选一，而是协同工作：

• GQA 减少 KV Cache 的存储量（注意力和缓存的存储都受益）

• KV Cache 减少每步的计算量（从 O(N²) 降到 O(N)）

• Flash Attention 减少每步的数据搬运量（避免注意力矩阵的 HBM 读写）

三者叠加的效果不是加法，而是乘法。以 LLaMA-3-8B 为例：使用 GQA-8 比 MHA 节省 75% 的 KV Cache；每步使用 KV Cache 加速约 4 倍；使用 Flash Attention 加速约 2-3 倍（长序列下）。三者结合让单张 A100-80GB 能够处理 65536 序列长度的推理，而标准 MHA + 无缓存 + 无 FA 的组合连 8192 都无法容纳，因为注意力矩阵和 KV Cache 的总和已经超过显存容量。

---

九、如何选择

9.1 决策树

训练还是推理？
├── 训练
│   ├── seq < 2048 → 标准 MHA（PyTorch 自动启用 FA）
│   ├── 2048-8192 → PyTorch F.sdp（自动选最优）
│   └── > 8192 → Flash Attention 或线性注意力
│
└── 推理
    ├── 需要长序列？→ GQA (n_groups=4或8) + FA + KV Cache
    ├── 显存受限？→ MQA 或 GQA-2（最大压缩比）
    └── 质量优先 → MHA（需更多显存）

9.2 实用建议

1. 训练阶段 ：用 PyTorch 2.0 的 F.scaled_dot_product_attention，自动启用 Flash Attention
2. 推理部署：使用 vLLM、TensorRT-LLM 或 llama.cpp，内置 GQA + KV Cache + PagedAttention
3. 微调适配：如果基础模型是 MHA（如 LLaMA-1），尝试 UPcycling 为 GQA 或直接使用 GQA 从头训练
4. 硬件选择：推理选高 HBM 带宽的 GPU（H100 的 3.35TB/s），比算力峰值更重要

9.3 与 DeepSeek 的结合方向

在实际应用中，DeepSeek 等大语言模型正是这些注意力优化技术的受益者。当你在 DeepSeek API 上处理长文档或构建 RAG 系统时：

1. 超长上下文：DeepSeek 最新模型支持高达 128K token 的上下文窗口，背后的关键技术之一就是 GQA 或 MLA 对 KV Cache 的优化。理解这些技术，你就能更好地理解 API 的使用限制和成本模型。

2. 推理部署：如果你在本地部署 DeepSeek 或其他大模型，GQA + Flash Attention + KV Cache 的配置直接决定了你的最大并发数和响应延迟。正确的注意力配置选择直接影响服务成本。

3. 性能调优：当处理长序列时（如代码库分析、长文档总结），注意瓶颈在数据搬运而非计算。在资源受限的环境中，GQA-2 或 MQA 可能是唯一可行的选择。

更多 DeepSeek 实战技巧，请参考：DeepSeek API 实战指南：从基础调用到高级应用

9.4 未来方向

注意力机制的研究仍在快速推进，以下是几个值得关注的趋势：

• MLA（Multi-head Latent Attention）：DeepSeek-V2 提出的创新方案。通过低秩压缩将 KV Cache 压缩为潜在向量，推理时再解压。在保持 MHA 质量的同时达到了比 MQA 更优的压缩率，是 GQA 路线的更激进版本。

• 线性注意力：Performer（随机 Fourier 特征近似）、Linformer（低秩投影）等方法将 O(N²) 降为 O(N)。质量有损，但适用于超长序列。

• 状态空间模型（SSM）：Mamba 系列模型用结构化状态空间模型替代注意力，实现了真正的线性序列复杂度。在长序列任务上展现出显著优势。

• 混合架构：Jamba、Samba 在部分层用注意力（精确 token 交互），其他层用 SSM（长期上下文传递）。不是非此即彼的选择，而是各取所长。

---

十、总结

本文从最基础的 Scaled Dot-Product Attention 出发，逐步深入到了实际大模型推理中至关重要的高级注意力机制。七年间，注意力机制从简单的矩阵乘法演进为了一个涵盖数学理论、GPU 硬件优化和系统设计的综合性领域。

回顾本文的核心收获：

✅ Scaled Dot-Product Attention：一切注意力的根基。QK^T / √d_k → Softmax → V 三步骤，每一步都有其必要性------矩阵乘法捕捉相关性，缩放解决高维梯度消失，Softmax 转化为概率，加权求和融合信息。

✅ MHA：将模型分裂成 h 个独立子空间，每个头关注不同的交互模式。标准配置（n_heads=12, d_k=64）沿用至今，但推理时的 KV Cache 显存开销随序列长度线性增长。

✅ MQA：保留多 Q 头、共享单 K/V 头的极端方案。KV Cache 减少 97%，推理速度提升 30-50%。质量损失在大模型中可接受。PaLM、Falcon 等模型证明了这条路线的可行性。

✅ GQA：MHA 和 MQA 的优雅折中，通过可配置的 n_groups 参数控制压缩比。在 n_groups=8 时几乎无质量损失，KV Cache 降低 75%。已成为 LLaMA 2/3 的标配，是目前最广泛采用的方案。

✅ KV Cache：推理阶段的核心优化，存储已计算的 K/V 避免冗余。与 GQA 配合，让单卡处理 128K 序列成为可能。

✅ Flash Attention：IO 感知的注意力算法。Tiling + Online Softmax 避免大矩阵 HBM 读写，在 A100 上实现 6-11x 加速。PyTorch 2.0 已内置，一行代码即可激活。

选择正确的注意力机制不是非黑即白的技术决策，它取决于你的场景和资源。但有一条原则可以指导选择：80% 的推理场景下，GQA-4 + Flash Attention + KV Cache 是最优配置------它不是任何一个维度的最优解，但在质量、速度、显存的三角权衡中给出了最好的综合表现。

---

本文代码已开源，所有内容原创。欢迎留言交流技术细节。

更多手写系列文章：DeepSeek 实战指南与完整教程汇总