引言如果让你用一句话概括过去七年人工智能领域最重要的技术突破答案几乎毫无悬念——注意力机制Attention Mechanism。2017年Google团队在论文《Attention Is All You Need》中首次提出Transformer架构彻底摒弃了当时主流的RNN和CNN结构仅靠注意力机制就刷新了机器翻译的SOTA记录。这篇论文的标题All You Need曾被很多人戏称为冥想指南但它却准确预言了深度学习发展的方向注意力机制不仅All You Need更是现代大语言模型的灵魂所在。八年后当我们回顾这场技术革命的全貌会发现注意力机制已经从最初的一种有效的技巧演变成了一套完整的工程体系——从标准的多头注意力到KV Cache优化MQA、GQA、MLA再到硬件感知的极致优化Flash Attention 1-4每一步演进都在解决真实问题也在为大模型的长上下文能力铺平道路。本文将系统梳理注意力机制的技术脉络从数学原理到工程实现从经典架构到最新进展带你深入理解这项支撑起大语言模型的核心技术。一、起源从Seq2Seq到注意力的出现1.1 Seq2Seq模型的困境在注意力机制出现之前序列到序列Seq2Seq模型是NLP领域的主流架构。典型的Seq2Seq模型由两部分组成编码器Encoder将输入序列编码为一个固定维度的上下文向量解码器Decoder基于这个上下文向量逐词生成输出序列这个架构看似合理却有一个致命的缺陷信息瓶颈。假设你正在将一篇5000字的文章翻译成另一种语言。编码器必须将整篇文章的信息压缩到一个几百维的向量中。这就像要求你把一整本书的内容背诵给朋友听朋友只能记住一个三行字的摘要——必然会有大量信息丢失。这个问题在处理长序列时尤为严重。随着序列增长模型性能急剧下降这就是著名的长距离依赖问题。1.2 注意力机制的诞生2015年Bahdanau等人提出了注意力机制的思想用于改善机器翻译质量。他们的核心洞察是不需要将整个输入序列压缩为一个向量而是在生成每个输出词时让模型关注输入序列中相关的部分。这个思想可以用一个直观的例子理解当你翻译The cat sat on the mat这句话时翻译到mat垫子这个词时你应该主要关注原文中对应的mat而不是平均地看所有词。注意力机制引入了三个关键概念查询Query当前位置想要查找的信息键Key每个位置的信息标签用于匹配查询值Value每个位置的实际内容通过Query和Key的匹配程度决定应该从Value中提取多少信息。二、核心原理Scaled Dot-Product Attention2.1 数学公式Transformer采用的是缩放点积注意力Scaled Dot-Product Attention其数学表达式为Attention(Q,K,V)softmax(dkQKT)V其中Q∈Rn×dk查询矩阵n为序列长度K∈Rm×dk键矩阵V∈Rm×dv值矩阵dk键向量的维度2.2 计算流程plaintext┌─────────────────────────────────────────────────────────┐ │ Scaled Dot-Product Attention 流程 │ ├─────────────────────────────────────────────────────────┤ │ │ │ 输入: Q, K, V 三组向量 │ │ │ │ Step 1: 计算相似度分数 │ │ ┌─────────┐ ┌─────────┐ │ │ │ Q │ × │ Kᵀ │ Score Matrix │ │ │ n×d_k │ │ d_k×m │ n×m │ │ └─────────┘ └─────────┘ │ │ │ │ Step 2: 缩放 │ │ ┌─────────────────────────┐ │ │ │ Score / √d_k │ 防止梯度消失/爆炸 │ │ └─────────────────────────┘ │ │ │ │ Step 3: Softmax归一化 │ │ ┌─────────────────────────┐ │ │ │ softmax(Score/√d_k) │ → Attention权重矩阵 │ │ └─────────────────────────┘ 每行和为1 │ │ │ │ Step 4: 加权求和 │ │ ┌─────────┐ ┌─────────┐ │ │ │权重矩阵 │ × │ V │ 输出向量 │ │ │ n×m │ │ m×d_v │ n×d_v │ │ └─────────┘ └─────────┘ │ │ │ └─────────────────────────────────────────────────────────┘2.3 为什么要缩放点积注意力的一个关键问题是方差膨胀。假设q和k是均值为0、方差为1的独立随机向量则它们的点积q⋅k的方差为dk。当dk较大时点积的值范围会很大导致softmax函数进入饱和区域梯度接近于零。除以dk可以将方差恢复到1确保softmax有足够大的梯度进行学习。2.4 PyTorch实现pythonimport torch import torch.nn as nn import torch.nn.functional as F import math class ScaledDotProductAttention(nn.Module): Scaled Dot-Product Attention 实现 公式: Attention(Q,K,V) softmax(QK^T / √d_k) * V def __init__(self, d_k: int, dropout: float 0.1): super().__init__() self.d_k d_k self.dropout nn.Dropout(dropout) def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor None) - tuple: 参数: q: [batch_size, n_heads, seq_len_q, d_k] k: [batch_size, n_heads, seq_len_k, d_k] v: [batch_size, n_heads, seq_len_k, d_v] mask: [batch_size, 1, seq_len_q, seq_len_k] 或 [batch_size, 1, 1, seq_len_k] 返回: output: [batch_size, n_heads, seq_len_q, d_v] attn_weights: [batch_size, n_heads, seq_len_q, seq_len_k] # 1. 计算点积相似度 [B, H, L_q, L_k] scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) # 2. 应用mask如有 if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) # 3. Softmax归一化得到注意力权重 attn_weights F.softmax(scores, dim-1) attn_weights self.dropout(attn_weights) # 4. 加权求和 output torch.matmul(attn_weights, v) return output, attn_weights三、多头注意力捕捉多维语义关系3.1 为什么需要多头单个注意力头只能捕捉一种类型的关系。语言是复杂的——有语法关系主谓宾、语义关系实体-属性、指代关系代词消解等。单一注意力头很难同时建模所有这些依赖。多头注意力Multi-Head Attention, MHA的核心思想是将Query、Key、Value投影到多个低维子空间在每个子空间中独立计算注意力最后拼接结果。3.2 数学公式MultiHead(Q,K,V)Concat(head1,…,headh)WO其中每个头的计算为headiAttention(QWiQ,KWiK,VWiV)WiQ,WiK,WiV∈Rdmodel×dk每头的投影矩阵WO∈Rhdv×dmodel输出投影矩阵h注意力头数通常设 dkdvdmodel/h3.3 多头注意力的直观理解plaintext┌──────────────────────────────────────────────────────────────────┐ │ Multi-Head Attention │ ├──────────────────────────────────────────────────────────────────┤ │ │ │ 输入 X (d_model) │ │ │ │ │ ▼ │ │ ┌─────────────────────────────────────────┐ │ │ │ Linear → 分 h 个头 │ │ │ └─────────────────────────────────────────┘ │ │ │ │ │ ├───────────────────────────────────┐ │ │ │ │ │ │ │ │ ▼ ▼ ▼ ▼ │ │ ┌────────┐ ┌────────┐ ┌────────┐ │ │ │ │ Head 1 │ │ Head 2 │ │ ... │ │ 各自独立计算 │ │ │ 关注 │ │ 关注 │ │ │ │ 注意力 │ │ │ 语法 │ │ 语义 │ │ │ │ │ │ └────┬───┘ └────┬───┘ └───┬───┘ │ │ │ │ │ │ │ │ │ └──────────────┼────────────┘ │ │ │ ▼ │ │ │ Concat (h × d_v) │ │ │ │ │ │ │ ▼ │ │ │ Linear → W^O │ │ │ │ │ │ │ ▼ │ │ │ 输出 (d_model) │ │ └──────────────────────────────────────────────────────────────────┘3.4 PyTorch实现pythonclass MultiHeadAttention(nn.Module): 多头注意力机制完整实现 特点: - 支持自注意力和交叉注意力 - 支持因果mask用于解码器 - KV Cache支持加速推理 def __init__(self, d_model: int, n_heads: int, dropout: float 0.1): super().__init__() assert d_model % n_heads 0, d_model 必须能被 n_heads 整除 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, biasFalse) self.W_k nn.Linear(d_model, d_model, biasFalse) self.W_v nn.Linear(d_model, d_model, biasFalse) self.W_o nn.Linear(d_model, d_model, biasFalse) self.dropout nn.Dropout(dropout) self.scaled_dot ScaledDotProductAttention(self.d_k, dropout) def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor None, kv_cache: dict None) - tuple: 参数: query: [batch_size, seq_len_q, d_model] key: [batch_size, seq_len_k, d_model] value: [batch_size, seq_len_v, d_model] mask: 注意力mask kv_cache: 可选的KV缓存 {k: ..., v: ...} 返回: output: [batch_size, seq_len_q, d_model] attn_weights: [batch_size, n_heads, seq_len_q, seq_len_k] batch_size query.size(0) # 1. 线性投影并分头 Q self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) K self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) V self.W_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) # 2. KV Cache处理 if kv_cache is not None and k in kv_cache: K torch.cat([kv_cache[k], K], dim2) V torch.cat([kv_cache[v], V], dim2) new_cache {k: K.detach(), v: V.detach()} # 3. 计算注意力 x, attn_weights self.scaled_dot(Q, K, V, mask) # 4. 合并多头并输出 x x.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) output self.W_o(x) return output, attn_weights, new_cache四、KV Cache优化注意力机制的效率革命4.1 问题注意力计算的二次复杂度标准注意力的核心问题是**O(n2)的时间和空间复杂度**。对于长度为n的序列注意力矩阵是n×n的——每个token都要与所有其他token交互。这在短序列上不是问题但当上下文扩展到32K、128K甚至更长时内存占用会爆炸式增长。更关键的是在推理阶段每次生成一个token都需要重新计算与所有历史token的注意力。如果不缓存中间结果每次生成都是O(n2)的计算灾难。4.2 KV Cache减少重复计算KV Cache是加速推理最基础的技术。其核心思想是在自回归生成过程中Key和Value对于已经处理过的token是不变的。只需要计算一次然后缓存起来。plaintext┌─────────────────────────────────────────────────────────────┐ │ KV Cache 工作原理 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ 生成第 t1 个token时 │ │ │ │ ┌─────────┐ ┌──────────────────────────────────┐ │ │ │ 历史KV │ │ 新Token Query │ │ │ │ Cache │ │ Q_{t1} │ │ │ │ K_{1:t} │ │ K_{t1}, V_{t1} (新计算) │ │ │ │ V_{1:t} │ │ │ │ │ └────┬────┘ └──────────┬───────────────────────┘ │ │ │ │ │ │ │ ┌──────────────┴───────────────┐ │ │ │ │ │ │ │ ▼ ▼ ▼ │ │ ┌──────────────┐ ┌──────────────┐ │ │ │ concat(K) │ × Q_{t1} → │ Attention │ │ │ │ concat(V) │ │ Score │ │ │ └──────────────┘ └──────────────┘ │ │ │ │ │ ▼ │ │ 生成 token_{t1} │ │ │ │ │ ▼ │ │ 更新 KV Cache │ │ │ └─────────────────────────────────────────────────────────────┘KV Cache显著加速了推理但也带来了新问题KV Cache的内存占用。4.3 从MHA到MQA减少KV头数Multi-Query AttentionMQA由Google在2019年提出其核心思想是所有Query头共享一组Key和Value。plaintext┌─────────────────────────────────────────────────────────────┐ │ MHA vs MQA 对比 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ MHA (Multi-Head Attention): │ │ Q: [n_heads × d_k] K: [n_heads × d_k] V: [n_heads × d_v] │ │ 每个头独立 K/V → 32头模型需要32组KV │ │ │ │ MQA (Multi-Query Attention): │ │ Q: [n_heads × d_k] K: [1 × d_k] V: [1 × d_v] │ │ 所有头共享1组K/V → 32头模型只需1组KV │ │ │ │ ┌────────────────────────────────────────────────────┐ │ │ │ 节省: (n_heads - 1) / n_heads 的KV Cache │ │ │ │ 例如: 32头模型可节省 ~97%的KV缓存空间 │ │ │ └────────────────────────────────────────────────────┘ │ │ │ └─────────────────────────────────────────────────────────────┘MQA的代价是表达能力下降。所有注意力头被迫使用相同的Key和Value难以捕捉多样化的关系。实验表明MQA可能导致模型质量下降5-10%。4.4 GQA寻找平衡点Grouped-Query AttentionGQA由Meta在Llama-2中提出是MHA和MQA的折中方案。核心思想将Query头分成若干组每组共享一组Key和Value。plaintext┌─────────────────────────────────────────────────────────────┐ │ MHA → GQA → MQA 演进 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ MHA: Q₁K₁V₁, Q₂K₂V₂, ..., Q₃₂K₃₂V₃₂ → 32组KV │ │ │ │ GQA: (Q₁Q₂)KV₁, (Q₃Q₄)KV₂, ..., → n_kv组KV │ │ (Q₃₁Q₃₂)KV₁₆ │ │ (以n_kv16为例) │ │ │ │ MQA: Q₁KV, Q₂KV, ..., Q₃₂KV → 1组KV │ │ │ │ ┌────────────────────────────────────────────────────┐ │ │ │ n_kv_heads 8 时: │ │ │ │ - KV Cache节省: 75% │ │ │ │ - 推理速度提升: ~1.6x │ │ │ │ - 质量损失: 3% │ │ │ └────────────────────────────────────────────────────┘ │ │ │ └─────────────────────────────────────────────────────────────┘GQA迅速成为主流方案Llama-3、Qwen2.5、Mistral、Gemma-2等主流开源模型都采用GQA。4.5 MLADeepSeek的创新突破Multi-head Latent AttentionMLA是DeepSeek在2024年提出的创新方案其核心思想更加精妙不共享KV而是对KV进行低秩压缩。plaintext┌─────────────────────────────────────────────────────────────┐ │ MLA低秩潜在注意力 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ 标准MHA: │ │ 缓存完整K: [d_heads × d_k] │ │ 缓存完整V: [d_heads × d_v] │ │ │ │ MLA: │ │ 1. 将K/V投影到低维空间 │ │ c_K W^{DK} X → [d_c] (d_c d_heads × d_k) │ │ c_V W^{DV} X → [d_c] │ │ │ │ 2. 缓存低维向量 c_K, c_V │ │ │ │ 3. 推理时用轻量MLP重建K/V │ │ K W^{UK} c_K │ │ V W^{UV} c_V │ │ │ │ ┌────────────────────────────────────────────────────┐ │ │ │ 效果对比基于DeepSeek-V2论文: │ │ │ │ - KV Cache显存降低: 5-10倍 │ │ │ │ - 支持128K上下文 │ │ │ │ - 建模质量: 与MHA相当甚至略优 │ │ │ └────────────────────────────────────────────────────┘ │ │ │ └─────────────────────────────────────────────────────────────┘MLA的关键洞察是不同注意力头的K/V之间存在大量冗余信息。所有头的K/V都源自同一个输入X底层的语义信息是共享的。MLA通过低秩分解捕获这种共享结构。4.6 四种机制对比表格指标MHAMQAGQA-8MLAKV Cache100%~2%~12%~8%推理速度1.0x2.1x1.6x1.8x建模质量100%92%98%100.5%适用场景训练/小模型极致压缩平衡之选长上下文五、长上下文优化稀疏注意力与Flash Attention5.1 稀疏注意力选择性关注标准注意力关注所有位置但直觉告诉我们并非所有位置都同等重要。稀疏注意力Sparse Attention的核心思想是每个token只与部分位置交互。常见的策略包括滑动窗口注意力Sliding Window每个token只关注周围w个token稀疏模式Fixed/Strided全局关注某些特殊位置如[CLS]局部全局混合结合局部模式捕捉细节全局模式捕捉整体Longformer和BigBird是稀疏注意力的代表模型它们在长文档处理上展现了优秀的能力。plaintext┌─────────────────────────────────────────────────────────────┐ │ 稀疏注意力模式示意 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ 标准注意力稠密: │ │ [● ● ● ● ● ● ● ● ● ●] │ │ ││││││││││ │ │ [● ● ● ● ● ● ● ● ● ●] 每个位置关注所有位置 │ │ │ │ 滑动窗口注意力: │ │ [● ● ● ● ● ● ● ● ● ●] │ │ │││││││ │ │ [ ● ● ● ● ● ● ● ● ●] 局部窗口w3 │ │ │ │ 全局滑动窗口: │ │ [● ● ●] ← 全局位置如[CLS] │ │ [ ● ● ● ● ● ● ● ● ●] ← 滑动窗口 │ │ │ └─────────────────────────────────────────────────────────────┘5.2 Flash Attention硬件感知的极致优化如果说稀疏注意力是算法层面的优化Flash Attention则是系统层面的革命。核心问题标准注意力需要在HBM高带宽内存中多次读写N×N的注意力矩阵。GPU的算力在增长但内存带宽增长缓慢导致内存访问成为瓶颈。Flash Attention的洞察GPU的SRAM比HBM快100倍以上但容量小得多。我们应该将注意力计算分块每次加载一块到SRAM中计算避免完整矩阵的HBM读写。Flash Attention v1 (2022)核心创新IO感知分块将N×N矩阵切分为适合SRAM的块在线softmax避免中间矩阵的存储核融合将整个注意力计算融合为一个GPU kernel效果内存复杂度从O(N2)降到O(N)速度提升最高7.6倍训练GPT-2提速3倍Flash Attention v2 (2023)改进更好的并行策略在序列长度维度并行化改进的工作分区提高GPU占用率效果A100利用率从25-40%提升到73%速度相比v1再提升1.7倍Flash Attention v3 (2024)Flash Attention-3专门为NVIDIAHopper架构H100优化利用了三个Hopper特性Tensor Memory Accelerator (TMA)专用硬件加载单元Warp Specialization生产者-消费者异步流水线FP8 Tensor Cores硬件支持的8位浮点运算plaintext┌─────────────────────────────────────────────────────────────┐ │ Flash Attention 演进对比 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ 版本 GPU 峰值算力 核心创新 │ │ ───────────────────────────────────────────────────────── │ │ FA1 A100 ~125 TFLOPS IO感知分块 │ │ FA2 A100 ~230 TFLOPS 更好的并行性 │ │ FA3 H100 ~740 TFLOPS Hoppe异步FP8 │ │ FA4 GB200 (ing) Blackwell原生优化 │ │ │ │ FA3在H100上实现75%利用率(vs FA2的35%) │ │ FP8精度接近1.2 PFLOPS精度损失可忽略 │ │ │ └─────────────────────────────────────────────────────────────┘5.3 Flash Attention的工程意义Flash Attention的重要性不仅在于速度更在于它让长上下文训练变得可能。标准注意力在长度为4K时注意力矩阵需要64MB长度为32K时需要4GB长度为128K时需要64GB。这远远超出了A100 80GB的显存限制。Flash Attention将显存需求降到O(N)使得在有限硬件上训练更长上下文成为可能。六、差分注意力最新研究方向2024年清华大学等机构提出了差分注意力Differential Attention机制这是注意力计算层面的创新。6.1 核心思想标准注意力计算单一的softmax分布。差分注意力的洞察是两个softmax分布的差可以消除噪声。DiffAttn(X)(softmax(dQ1K1T)−λ⋅softmax(dQ2K2T))V这个设计类似于差分放大器或降噪耳机共同模式的信号被抑制真正有用的信号被放大。6.2 实验效果在3B规模模型上的实验显示幻觉减少TruthfulQA上下文学习能力提升长上下文检索更鲁棒激活异常值减少有利于量化差分注意力的意义是概念性的它表明softmax分布本身还有未被挖掘的结构——注意力不仅仅是稀疏的还是有噪声的而噪声是部分系统性且可消除的。七、注意力机制在各模型架构中的应用7.1 主流模型注意力配置表格模型注意力类型上下文长度特点GPT-4MHA128K每个头独立KVLlama-3GQA-8128K8组KVQwen-2.5GQA128K平衡效率DeepSeek-V3MLA128K低秩压缩Mistral-7BGQA32KSliding Window7.2 DeepSeek-V3的MLA实现DeepSeek-V3采用MLA的核心挑战是训练时需要存储完整的上投影矩阵。DeepSeek通过巧妙的设计解决了这一问题python# MLA的PyTorch伪代码 class MultiHeadLatentAttention(nn.Module): def __init__(self, d_model, n_heads, d_c): super().__init__() # 低秩压缩矩阵推理时使用训练时需保留完整矩阵 self.W_DK nn.Linear(d_model, d_c) # 下投影 self.W_UK nn.Linear(d_c, d_model) # K的上投影 self.W_DV nn.Linear(d_model, d_c) # V的下投影 self.W_UV nn.Linear(d_c, d_model) # V的上投影 # Query保持标准形式 self.W_Q nn.Linear(d_model, d_model) def forward(self, x): # 低秩压缩 c_K self.W_DK(x) # [B, L, d_c] c_V self.W_DV(x) # [B, L, d_c] # 重建K/V推理时缓存c_K, c_V而非K, V K self.W_UK(c_K) # [B, L, d_model] V self.W_UV(c_V) # [B, L, d_model] # 标准注意力计算 Q self.W_Q(x) # ... 后续计算7.3 混合架构注意力与状态空间模型的结合Mamba等状态空间模型SSM提供了注意力的替代方案O(n)的时间复杂度和常数级内存。混合架构正在成为新趋势Jamba注意力层和Mamba层按1:7比例交织Nemotron类似的混合设计这表明注意力和SSM是互补而非竞争的——注意力擅长精确的位置检索SSM擅长高效处理序列模式。八、未来展望与思考8.1 当前挑战长上下文的信息瓶颈即使有Flash Attention模型在超长上下文中的表现仍有下降注意力分数的噪声差分注意力等研究正在探索训练与推理的不对称MLA等方案在推理阶段优化但训练成本仍然很高注意力机制的脆弱性最新研究显示在时间序列等任务中注意力可能退化为MLP8.2 未来方向更高效的稀疏模式可学习的稀疏注意力持续内存机制结合外部记忆的注意力硬件-算法协同设计Flash Attention 4针对Blackwell架构的优化混合架构的深化注意力与SSM、更高效RNN变体的结合8.3 笔者的思考回顾注意力机制的演进历程有一个清晰的脉络每一代技术突破都在解决前一代的核心矛盾。MHA解决了表达力问题但带来了内存问题MQA/GQA解决了内存问题但牺牲了部分表达力MLA试图同时解决两者通过低秩结构找到新的平衡Flash Attention则从另一个维度突破——不改变算法而是让硬件更高效这个演进告诉我们没有完美的架构只有当前条件下的最优解。理解每种技术的取舍才能在大模型实践中做出正确的选择。九、总结本文系统梳理了大模型注意力机制的技术全貌原理层面从生物启发的注意力思想到Scaled Dot-Product Attention的数学形式化再到多头注意力捕捉多维语义关系效率层面从MQA、GQA到MLA的KV Cache压缩技术每种方案都在效率与质量之间寻找平衡工程层面Flash Attention开创了硬件感知算法设计的先河将注意力计算从O(N2)内存降到O(N)架构层面不同模型选择了不同的注意力方案反映了各自场景的需求和权衡注意力机制的故事远未结束。随着上下文窗口扩展到百万级多模态成为标配我们期待看到更多创新。但无论如何演进让模型学会关注重要信息这个核心思想将始终是人工智能最重要的洞见之一。参考资料Vaswani et al., Attention Is All You Need, NeurIPS 2017Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023Shah et al., FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision, 2024Ainslie et al., GQA: Training Generalized Multi-Query Transformer, 2023DeepSeek-V2 Technical Report, 2024Ye et al., Differential Transformer, 2024本文约5500字适合技术从业者深度阅读。如有问题或讨论欢迎在评论区交流。
大模型注意力机制深度解析:从Dot-Product到Flash Attention的演进之路
发布时间:2026/5/20 1:49:07
引言如果让你用一句话概括过去七年人工智能领域最重要的技术突破答案几乎毫无悬念——注意力机制Attention Mechanism。2017年Google团队在论文《Attention Is All You Need》中首次提出Transformer架构彻底摒弃了当时主流的RNN和CNN结构仅靠注意力机制就刷新了机器翻译的SOTA记录。这篇论文的标题All You Need曾被很多人戏称为冥想指南但它却准确预言了深度学习发展的方向注意力机制不仅All You Need更是现代大语言模型的灵魂所在。八年后当我们回顾这场技术革命的全貌会发现注意力机制已经从最初的一种有效的技巧演变成了一套完整的工程体系——从标准的多头注意力到KV Cache优化MQA、GQA、MLA再到硬件感知的极致优化Flash Attention 1-4每一步演进都在解决真实问题也在为大模型的长上下文能力铺平道路。本文将系统梳理注意力机制的技术脉络从数学原理到工程实现从经典架构到最新进展带你深入理解这项支撑起大语言模型的核心技术。一、起源从Seq2Seq到注意力的出现1.1 Seq2Seq模型的困境在注意力机制出现之前序列到序列Seq2Seq模型是NLP领域的主流架构。典型的Seq2Seq模型由两部分组成编码器Encoder将输入序列编码为一个固定维度的上下文向量解码器Decoder基于这个上下文向量逐词生成输出序列这个架构看似合理却有一个致命的缺陷信息瓶颈。假设你正在将一篇5000字的文章翻译成另一种语言。编码器必须将整篇文章的信息压缩到一个几百维的向量中。这就像要求你把一整本书的内容背诵给朋友听朋友只能记住一个三行字的摘要——必然会有大量信息丢失。这个问题在处理长序列时尤为严重。随着序列增长模型性能急剧下降这就是著名的长距离依赖问题。1.2 注意力机制的诞生2015年Bahdanau等人提出了注意力机制的思想用于改善机器翻译质量。他们的核心洞察是不需要将整个输入序列压缩为一个向量而是在生成每个输出词时让模型关注输入序列中相关的部分。这个思想可以用一个直观的例子理解当你翻译The cat sat on the mat这句话时翻译到mat垫子这个词时你应该主要关注原文中对应的mat而不是平均地看所有词。注意力机制引入了三个关键概念查询Query当前位置想要查找的信息键Key每个位置的信息标签用于匹配查询值Value每个位置的实际内容通过Query和Key的匹配程度决定应该从Value中提取多少信息。二、核心原理Scaled Dot-Product Attention2.1 数学公式Transformer采用的是缩放点积注意力Scaled Dot-Product Attention其数学表达式为Attention(Q,K,V)softmax(dkQKT)V其中Q∈Rn×dk查询矩阵n为序列长度K∈Rm×dk键矩阵V∈Rm×dv值矩阵dk键向量的维度2.2 计算流程plaintext┌─────────────────────────────────────────────────────────┐ │ Scaled Dot-Product Attention 流程 │ ├─────────────────────────────────────────────────────────┤ │ │ │ 输入: Q, K, V 三组向量 │ │ │ │ Step 1: 计算相似度分数 │ │ ┌─────────┐ ┌─────────┐ │ │ │ Q │ × │ Kᵀ │ Score Matrix │ │ │ n×d_k │ │ d_k×m │ n×m │ │ └─────────┘ └─────────┘ │ │ │ │ Step 2: 缩放 │ │ ┌─────────────────────────┐ │ │ │ Score / √d_k │ 防止梯度消失/爆炸 │ │ └─────────────────────────┘ │ │ │ │ Step 3: Softmax归一化 │ │ ┌─────────────────────────┐ │ │ │ softmax(Score/√d_k) │ → Attention权重矩阵 │ │ └─────────────────────────┘ 每行和为1 │ │ │ │ Step 4: 加权求和 │ │ ┌─────────┐ ┌─────────┐ │ │ │权重矩阵 │ × │ V │ 输出向量 │ │ │ n×m │ │ m×d_v │ n×d_v │ │ └─────────┘ └─────────┘ │ │ │ └─────────────────────────────────────────────────────────┘2.3 为什么要缩放点积注意力的一个关键问题是方差膨胀。假设q和k是均值为0、方差为1的独立随机向量则它们的点积q⋅k的方差为dk。当dk较大时点积的值范围会很大导致softmax函数进入饱和区域梯度接近于零。除以dk可以将方差恢复到1确保softmax有足够大的梯度进行学习。2.4 PyTorch实现pythonimport torch import torch.nn as nn import torch.nn.functional as F import math class ScaledDotProductAttention(nn.Module): Scaled Dot-Product Attention 实现 公式: Attention(Q,K,V) softmax(QK^T / √d_k) * V def __init__(self, d_k: int, dropout: float 0.1): super().__init__() self.d_k d_k self.dropout nn.Dropout(dropout) def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor None) - tuple: 参数: q: [batch_size, n_heads, seq_len_q, d_k] k: [batch_size, n_heads, seq_len_k, d_k] v: [batch_size, n_heads, seq_len_k, d_v] mask: [batch_size, 1, seq_len_q, seq_len_k] 或 [batch_size, 1, 1, seq_len_k] 返回: output: [batch_size, n_heads, seq_len_q, d_v] attn_weights: [batch_size, n_heads, seq_len_q, seq_len_k] # 1. 计算点积相似度 [B, H, L_q, L_k] scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) # 2. 应用mask如有 if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) # 3. Softmax归一化得到注意力权重 attn_weights F.softmax(scores, dim-1) attn_weights self.dropout(attn_weights) # 4. 加权求和 output torch.matmul(attn_weights, v) return output, attn_weights三、多头注意力捕捉多维语义关系3.1 为什么需要多头单个注意力头只能捕捉一种类型的关系。语言是复杂的——有语法关系主谓宾、语义关系实体-属性、指代关系代词消解等。单一注意力头很难同时建模所有这些依赖。多头注意力Multi-Head Attention, MHA的核心思想是将Query、Key、Value投影到多个低维子空间在每个子空间中独立计算注意力最后拼接结果。3.2 数学公式MultiHead(Q,K,V)Concat(head1,…,headh)WO其中每个头的计算为headiAttention(QWiQ,KWiK,VWiV)WiQ,WiK,WiV∈Rdmodel×dk每头的投影矩阵WO∈Rhdv×dmodel输出投影矩阵h注意力头数通常设 dkdvdmodel/h3.3 多头注意力的直观理解plaintext┌──────────────────────────────────────────────────────────────────┐ │ Multi-Head Attention │ ├──────────────────────────────────────────────────────────────────┤ │ │ │ 输入 X (d_model) │ │ │ │ │ ▼ │ │ ┌─────────────────────────────────────────┐ │ │ │ Linear → 分 h 个头 │ │ │ └─────────────────────────────────────────┘ │ │ │ │ │ ├───────────────────────────────────┐ │ │ │ │ │ │ │ │ ▼ ▼ ▼ ▼ │ │ ┌────────┐ ┌────────┐ ┌────────┐ │ │ │ │ Head 1 │ │ Head 2 │ │ ... │ │ 各自独立计算 │ │ │ 关注 │ │ 关注 │ │ │ │ 注意力 │ │ │ 语法 │ │ 语义 │ │ │ │ │ │ └────┬───┘ └────┬───┘ └───┬───┘ │ │ │ │ │ │ │ │ │ └──────────────┼────────────┘ │ │ │ ▼ │ │ │ Concat (h × d_v) │ │ │ │ │ │ │ ▼ │ │ │ Linear → W^O │ │ │ │ │ │ │ ▼ │ │ │ 输出 (d_model) │ │ └──────────────────────────────────────────────────────────────────┘3.4 PyTorch实现pythonclass MultiHeadAttention(nn.Module): 多头注意力机制完整实现 特点: - 支持自注意力和交叉注意力 - 支持因果mask用于解码器 - KV Cache支持加速推理 def __init__(self, d_model: int, n_heads: int, dropout: float 0.1): super().__init__() assert d_model % n_heads 0, d_model 必须能被 n_heads 整除 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, biasFalse) self.W_k nn.Linear(d_model, d_model, biasFalse) self.W_v nn.Linear(d_model, d_model, biasFalse) self.W_o nn.Linear(d_model, d_model, biasFalse) self.dropout nn.Dropout(dropout) self.scaled_dot ScaledDotProductAttention(self.d_k, dropout) def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor None, kv_cache: dict None) - tuple: 参数: query: [batch_size, seq_len_q, d_model] key: [batch_size, seq_len_k, d_model] value: [batch_size, seq_len_v, d_model] mask: 注意力mask kv_cache: 可选的KV缓存 {k: ..., v: ...} 返回: output: [batch_size, seq_len_q, d_model] attn_weights: [batch_size, n_heads, seq_len_q, seq_len_k] batch_size query.size(0) # 1. 线性投影并分头 Q self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) K self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) V self.W_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) # 2. KV Cache处理 if kv_cache is not None and k in kv_cache: K torch.cat([kv_cache[k], K], dim2) V torch.cat([kv_cache[v], V], dim2) new_cache {k: K.detach(), v: V.detach()} # 3. 计算注意力 x, attn_weights self.scaled_dot(Q, K, V, mask) # 4. 合并多头并输出 x x.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) output self.W_o(x) return output, attn_weights, new_cache四、KV Cache优化注意力机制的效率革命4.1 问题注意力计算的二次复杂度标准注意力的核心问题是**O(n2)的时间和空间复杂度**。对于长度为n的序列注意力矩阵是n×n的——每个token都要与所有其他token交互。这在短序列上不是问题但当上下文扩展到32K、128K甚至更长时内存占用会爆炸式增长。更关键的是在推理阶段每次生成一个token都需要重新计算与所有历史token的注意力。如果不缓存中间结果每次生成都是O(n2)的计算灾难。4.2 KV Cache减少重复计算KV Cache是加速推理最基础的技术。其核心思想是在自回归生成过程中Key和Value对于已经处理过的token是不变的。只需要计算一次然后缓存起来。plaintext┌─────────────────────────────────────────────────────────────┐ │ KV Cache 工作原理 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ 生成第 t1 个token时 │ │ │ │ ┌─────────┐ ┌──────────────────────────────────┐ │ │ │ 历史KV │ │ 新Token Query │ │ │ │ Cache │ │ Q_{t1} │ │ │ │ K_{1:t} │ │ K_{t1}, V_{t1} (新计算) │ │ │ │ V_{1:t} │ │ │ │ │ └────┬────┘ └──────────┬───────────────────────┘ │ │ │ │ │ │ │ ┌──────────────┴───────────────┐ │ │ │ │ │ │ │ ▼ ▼ ▼ │ │ ┌──────────────┐ ┌──────────────┐ │ │ │ concat(K) │ × Q_{t1} → │ Attention │ │ │ │ concat(V) │ │ Score │ │ │ └──────────────┘ └──────────────┘ │ │ │ │ │ ▼ │ │ 生成 token_{t1} │ │ │ │ │ ▼ │ │ 更新 KV Cache │ │ │ └─────────────────────────────────────────────────────────────┘KV Cache显著加速了推理但也带来了新问题KV Cache的内存占用。4.3 从MHA到MQA减少KV头数Multi-Query AttentionMQA由Google在2019年提出其核心思想是所有Query头共享一组Key和Value。plaintext┌─────────────────────────────────────────────────────────────┐ │ MHA vs MQA 对比 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ MHA (Multi-Head Attention): │ │ Q: [n_heads × d_k] K: [n_heads × d_k] V: [n_heads × d_v] │ │ 每个头独立 K/V → 32头模型需要32组KV │ │ │ │ MQA (Multi-Query Attention): │ │ Q: [n_heads × d_k] K: [1 × d_k] V: [1 × d_v] │ │ 所有头共享1组K/V → 32头模型只需1组KV │ │ │ │ ┌────────────────────────────────────────────────────┐ │ │ │ 节省: (n_heads - 1) / n_heads 的KV Cache │ │ │ │ 例如: 32头模型可节省 ~97%的KV缓存空间 │ │ │ └────────────────────────────────────────────────────┘ │ │ │ └─────────────────────────────────────────────────────────────┘MQA的代价是表达能力下降。所有注意力头被迫使用相同的Key和Value难以捕捉多样化的关系。实验表明MQA可能导致模型质量下降5-10%。4.4 GQA寻找平衡点Grouped-Query AttentionGQA由Meta在Llama-2中提出是MHA和MQA的折中方案。核心思想将Query头分成若干组每组共享一组Key和Value。plaintext┌─────────────────────────────────────────────────────────────┐ │ MHA → GQA → MQA 演进 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ MHA: Q₁K₁V₁, Q₂K₂V₂, ..., Q₃₂K₃₂V₃₂ → 32组KV │ │ │ │ GQA: (Q₁Q₂)KV₁, (Q₃Q₄)KV₂, ..., → n_kv组KV │ │ (Q₃₁Q₃₂)KV₁₆ │ │ (以n_kv16为例) │ │ │ │ MQA: Q₁KV, Q₂KV, ..., Q₃₂KV → 1组KV │ │ │ │ ┌────────────────────────────────────────────────────┐ │ │ │ n_kv_heads 8 时: │ │ │ │ - KV Cache节省: 75% │ │ │ │ - 推理速度提升: ~1.6x │ │ │ │ - 质量损失: 3% │ │ │ └────────────────────────────────────────────────────┘ │ │ │ └─────────────────────────────────────────────────────────────┘GQA迅速成为主流方案Llama-3、Qwen2.5、Mistral、Gemma-2等主流开源模型都采用GQA。4.5 MLADeepSeek的创新突破Multi-head Latent AttentionMLA是DeepSeek在2024年提出的创新方案其核心思想更加精妙不共享KV而是对KV进行低秩压缩。plaintext┌─────────────────────────────────────────────────────────────┐ │ MLA低秩潜在注意力 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ 标准MHA: │ │ 缓存完整K: [d_heads × d_k] │ │ 缓存完整V: [d_heads × d_v] │ │ │ │ MLA: │ │ 1. 将K/V投影到低维空间 │ │ c_K W^{DK} X → [d_c] (d_c d_heads × d_k) │ │ c_V W^{DV} X → [d_c] │ │ │ │ 2. 缓存低维向量 c_K, c_V │ │ │ │ 3. 推理时用轻量MLP重建K/V │ │ K W^{UK} c_K │ │ V W^{UV} c_V │ │ │ │ ┌────────────────────────────────────────────────────┐ │ │ │ 效果对比基于DeepSeek-V2论文: │ │ │ │ - KV Cache显存降低: 5-10倍 │ │ │ │ - 支持128K上下文 │ │ │ │ - 建模质量: 与MHA相当甚至略优 │ │ │ └────────────────────────────────────────────────────┘ │ │ │ └─────────────────────────────────────────────────────────────┘MLA的关键洞察是不同注意力头的K/V之间存在大量冗余信息。所有头的K/V都源自同一个输入X底层的语义信息是共享的。MLA通过低秩分解捕获这种共享结构。4.6 四种机制对比表格指标MHAMQAGQA-8MLAKV Cache100%~2%~12%~8%推理速度1.0x2.1x1.6x1.8x建模质量100%92%98%100.5%适用场景训练/小模型极致压缩平衡之选长上下文五、长上下文优化稀疏注意力与Flash Attention5.1 稀疏注意力选择性关注标准注意力关注所有位置但直觉告诉我们并非所有位置都同等重要。稀疏注意力Sparse Attention的核心思想是每个token只与部分位置交互。常见的策略包括滑动窗口注意力Sliding Window每个token只关注周围w个token稀疏模式Fixed/Strided全局关注某些特殊位置如[CLS]局部全局混合结合局部模式捕捉细节全局模式捕捉整体Longformer和BigBird是稀疏注意力的代表模型它们在长文档处理上展现了优秀的能力。plaintext┌─────────────────────────────────────────────────────────────┐ │ 稀疏注意力模式示意 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ 标准注意力稠密: │ │ [● ● ● ● ● ● ● ● ● ●] │ │ ││││││││││ │ │ [● ● ● ● ● ● ● ● ● ●] 每个位置关注所有位置 │ │ │ │ 滑动窗口注意力: │ │ [● ● ● ● ● ● ● ● ● ●] │ │ │││││││ │ │ [ ● ● ● ● ● ● ● ● ●] 局部窗口w3 │ │ │ │ 全局滑动窗口: │ │ [● ● ●] ← 全局位置如[CLS] │ │ [ ● ● ● ● ● ● ● ● ●] ← 滑动窗口 │ │ │ └─────────────────────────────────────────────────────────────┘5.2 Flash Attention硬件感知的极致优化如果说稀疏注意力是算法层面的优化Flash Attention则是系统层面的革命。核心问题标准注意力需要在HBM高带宽内存中多次读写N×N的注意力矩阵。GPU的算力在增长但内存带宽增长缓慢导致内存访问成为瓶颈。Flash Attention的洞察GPU的SRAM比HBM快100倍以上但容量小得多。我们应该将注意力计算分块每次加载一块到SRAM中计算避免完整矩阵的HBM读写。Flash Attention v1 (2022)核心创新IO感知分块将N×N矩阵切分为适合SRAM的块在线softmax避免中间矩阵的存储核融合将整个注意力计算融合为一个GPU kernel效果内存复杂度从O(N2)降到O(N)速度提升最高7.6倍训练GPT-2提速3倍Flash Attention v2 (2023)改进更好的并行策略在序列长度维度并行化改进的工作分区提高GPU占用率效果A100利用率从25-40%提升到73%速度相比v1再提升1.7倍Flash Attention v3 (2024)Flash Attention-3专门为NVIDIAHopper架构H100优化利用了三个Hopper特性Tensor Memory Accelerator (TMA)专用硬件加载单元Warp Specialization生产者-消费者异步流水线FP8 Tensor Cores硬件支持的8位浮点运算plaintext┌─────────────────────────────────────────────────────────────┐ │ Flash Attention 演进对比 │ ├─────────────────────────────────────────────────────────────┤ │ │ │ 版本 GPU 峰值算力 核心创新 │ │ ───────────────────────────────────────────────────────── │ │ FA1 A100 ~125 TFLOPS IO感知分块 │ │ FA2 A100 ~230 TFLOPS 更好的并行性 │ │ FA3 H100 ~740 TFLOPS Hoppe异步FP8 │ │ FA4 GB200 (ing) Blackwell原生优化 │ │ │ │ FA3在H100上实现75%利用率(vs FA2的35%) │ │ FP8精度接近1.2 PFLOPS精度损失可忽略 │ │ │ └─────────────────────────────────────────────────────────────┘5.3 Flash Attention的工程意义Flash Attention的重要性不仅在于速度更在于它让长上下文训练变得可能。标准注意力在长度为4K时注意力矩阵需要64MB长度为32K时需要4GB长度为128K时需要64GB。这远远超出了A100 80GB的显存限制。Flash Attention将显存需求降到O(N)使得在有限硬件上训练更长上下文成为可能。六、差分注意力最新研究方向2024年清华大学等机构提出了差分注意力Differential Attention机制这是注意力计算层面的创新。6.1 核心思想标准注意力计算单一的softmax分布。差分注意力的洞察是两个softmax分布的差可以消除噪声。DiffAttn(X)(softmax(dQ1K1T)−λ⋅softmax(dQ2K2T))V这个设计类似于差分放大器或降噪耳机共同模式的信号被抑制真正有用的信号被放大。6.2 实验效果在3B规模模型上的实验显示幻觉减少TruthfulQA上下文学习能力提升长上下文检索更鲁棒激活异常值减少有利于量化差分注意力的意义是概念性的它表明softmax分布本身还有未被挖掘的结构——注意力不仅仅是稀疏的还是有噪声的而噪声是部分系统性且可消除的。七、注意力机制在各模型架构中的应用7.1 主流模型注意力配置表格模型注意力类型上下文长度特点GPT-4MHA128K每个头独立KVLlama-3GQA-8128K8组KVQwen-2.5GQA128K平衡效率DeepSeek-V3MLA128K低秩压缩Mistral-7BGQA32KSliding Window7.2 DeepSeek-V3的MLA实现DeepSeek-V3采用MLA的核心挑战是训练时需要存储完整的上投影矩阵。DeepSeek通过巧妙的设计解决了这一问题python# MLA的PyTorch伪代码 class MultiHeadLatentAttention(nn.Module): def __init__(self, d_model, n_heads, d_c): super().__init__() # 低秩压缩矩阵推理时使用训练时需保留完整矩阵 self.W_DK nn.Linear(d_model, d_c) # 下投影 self.W_UK nn.Linear(d_c, d_model) # K的上投影 self.W_DV nn.Linear(d_model, d_c) # V的下投影 self.W_UV nn.Linear(d_c, d_model) # V的上投影 # Query保持标准形式 self.W_Q nn.Linear(d_model, d_model) def forward(self, x): # 低秩压缩 c_K self.W_DK(x) # [B, L, d_c] c_V self.W_DV(x) # [B, L, d_c] # 重建K/V推理时缓存c_K, c_V而非K, V K self.W_UK(c_K) # [B, L, d_model] V self.W_UV(c_V) # [B, L, d_model] # 标准注意力计算 Q self.W_Q(x) # ... 后续计算7.3 混合架构注意力与状态空间模型的结合Mamba等状态空间模型SSM提供了注意力的替代方案O(n)的时间复杂度和常数级内存。混合架构正在成为新趋势Jamba注意力层和Mamba层按1:7比例交织Nemotron类似的混合设计这表明注意力和SSM是互补而非竞争的——注意力擅长精确的位置检索SSM擅长高效处理序列模式。八、未来展望与思考8.1 当前挑战长上下文的信息瓶颈即使有Flash Attention模型在超长上下文中的表现仍有下降注意力分数的噪声差分注意力等研究正在探索训练与推理的不对称MLA等方案在推理阶段优化但训练成本仍然很高注意力机制的脆弱性最新研究显示在时间序列等任务中注意力可能退化为MLP8.2 未来方向更高效的稀疏模式可学习的稀疏注意力持续内存机制结合外部记忆的注意力硬件-算法协同设计Flash Attention 4针对Blackwell架构的优化混合架构的深化注意力与SSM、更高效RNN变体的结合8.3 笔者的思考回顾注意力机制的演进历程有一个清晰的脉络每一代技术突破都在解决前一代的核心矛盾。MHA解决了表达力问题但带来了内存问题MQA/GQA解决了内存问题但牺牲了部分表达力MLA试图同时解决两者通过低秩结构找到新的平衡Flash Attention则从另一个维度突破——不改变算法而是让硬件更高效这个演进告诉我们没有完美的架构只有当前条件下的最优解。理解每种技术的取舍才能在大模型实践中做出正确的选择。九、总结本文系统梳理了大模型注意力机制的技术全貌原理层面从生物启发的注意力思想到Scaled Dot-Product Attention的数学形式化再到多头注意力捕捉多维语义关系效率层面从MQA、GQA到MLA的KV Cache压缩技术每种方案都在效率与质量之间寻找平衡工程层面Flash Attention开创了硬件感知算法设计的先河将注意力计算从O(N2)内存降到O(N)架构层面不同模型选择了不同的注意力方案反映了各自场景的需求和权衡注意力机制的故事远未结束。随着上下文窗口扩展到百万级多模态成为标配我们期待看到更多创新。但无论如何演进让模型学会关注重要信息这个核心思想将始终是人工智能最重要的洞见之一。参考资料Vaswani et al., Attention Is All You Need, NeurIPS 2017Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023Shah et al., FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision, 2024Ainslie et al., GQA: Training Generalized Multi-Query Transformer, 2023DeepSeek-V2 Technical Report, 2024Ye et al., Differential Transformer, 2024本文约5500字适合技术从业者深度阅读。如有问题或讨论欢迎在评论区交流。