ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

从 MLA 到 CSA,聊聊大模型 Attention 的 “瘦身” 与 “闪送”

从 MLA 到 CSA,聊聊大模型 Attention 的 “瘦身” 与 “闪送” 最近做了太多应用层的事底层原理反而搁置了。趁着放假围绕 Attention 机制聊聊这两年注意力架构的技术演化路线。全文覆盖三个方向每个方向都是上一个方向的接力棒KV Cache 压缩类MQA → GQA → MLA每个 token 的 KV 变细推理效率优化类Flash Attention、Paged Attention让 GPU 跑得更满长文本优化类NSA → DSA → CSA HCA需要 attend 的 token 变少一、MHA从古老的Multi Head Attention说起哈哈LLM 技术一天往往人间一年。说是古老掐指一算Transformer 论文正式发表也不过才 10 年不到。但就在这短短几年里MHA 已经从标准答案变成了需要优化的起点。标准 MHAMulti-Head Attention的逻辑很直接为每个注意力头独立计算 Q、K、V 投影每个 token 都和所有历史 token 建立关联。问题也很直接问题说明KV Cache内存爆炸长上下文128K或较大batch时KV Cache迅速成为GPU显存瓶颈推理带宽受限每生成一个新token需要从HBM加载整个历史KV Cache到SRAM进行attention计算带宽需求随heads和seq_len线性增长长上下文计算代价高估算在解码 64k 长度上下文时基于 softmax 的 attention 计算占总延迟的 70–80%凸显出对更高效 attention 机制的迫切需求核心矛盾只有一个KV 太胖历史太长搬不动。优化方向也随之分化成两条路要么让每个 token 的 KV 变瘦纵向压缩要么让需要 attend 的 token 变少横向稀疏。二、压缩 KV Cache让每个 token 的 KV 变瘦2.1 MQA激进共享核心设计多个 Query head只保留一个 K head 和一个 V head。概念就是参数共享MQAMulti-Query Attention是第一个掀桌子的方案。如果 32 个 Query head 都用同一套 KVKV Cache 直接缩小 32 倍。代价是什么精度有损失。把32 套独立的视角强行合并成1 套公共视角不同 head 之间的多样性丢失了。所以 MQA 一般只用在对速度极度敏感、容忍一定精度下降的场景。2.2 GQA折中方案 ⭐核心设计引入差值概念提供介于MQA和MHA之间在质量与速度之间提供了良好的折中案GQAGrouped Query Attention引入了一个参数 G分组数G num_heads退化为 MHA每个 Q head 独享 KVG 1退化为 MQA所有 Q head 共用一套 KV1 G num_headsGQA 的甜蜜区间质量与速度的折中主流模型的压缩幅度一般在4~8倍def GQA(X, W_Q, W_K, W_V, W_O, num_q_heads, num_kv_heads, d_k): num_q_heads: Query 头数 (e.g., 32) num_kv_heads: KV 头数 (e.g., 8)即分组数 G 每组 num_q_heads//num_kv_heads 个 Q 共享一对 K/V B, T, D X.shape groups num_q_heads // num_kv_heads # 每个KV head服务的Q head数 # Q 有 num_q_heads 个 Q (X W_Q).view(B, T, num_q_heads, d_k).transpose(1, 2) # [B, Hq, T, d_k] # K/V 只有 num_kv_heads 个 K (X W_K).view(B, T, num_kv_heads, d_k).transpose(1, 2) # [B, Hkv, T, d_k] V (X W_V).view(B, T, num_kv_heads, d_k).transpose(1, 2) # [B, Hkv, T, d_k] # 将 K/V 重复扩展以匹配 Q 的头数 (或用 reshape 广播) K K.repeat_interleave(groups, dim1) # [B, Hq, T, d_k] V V.repeat_interleave(groups, dim1) scores (Q K.transpose(-2, -1)) / sqrt(d_k) # [B, Hq, T, T] attn softmax(scores, dim-1) out attn V # [B, Hq, T, d_k] out out.transpose(1, 2).contiguous().view(B, T, D) return out W_O2.3 MLA低秩压缩DeepSeek 的独门绝技 ⭐核心设计不通过参数共享而是通过矩阵低秩压缩来缩小 KV。说到低秩压缩熟悉微调的同学一定不陌生——LoRA 背后就是同一个思想。MLA引入了两个投影矩阵通过在存储KV时对KV进行下投影压缩到更小的矩阵存储再在推理时通过上投影矩阵恢复到原有维度来降低KV Cache的存储量级。但 MLA 面临一个额外的挑战KV 里有 RoPE 位置编码。位置编码必须感知序列顺序不能被低秩压缩抹平。所以 MLA 把 K 拆成两部分NoPE 部分无位置编码KV共同走低秩压缩通道存入压缩潜变量 c_KVRoPE 部分有位置编码QK各保留一部分维度经过ROPE处理以保持对位置变量的敏感性不压缩单独保留最终推理时通过把RoPE和NoPE进行拼接做Attention运算。def MLA(X, params, num_heads, d_model, d_c, d_r): d_c: KV 压缩后的潜在维度 (远小于 num_heads * d_k) d_r: RoPE 分量维度 B, T, D X.shape # Step 1: KV 压缩核心 # 下投影将 X 压缩为低秩潜变量 c_kv c_kv X params[W_DKV] # [B, T, d_c] d_c num_heads*d_k # 缓存 c_kv 即可不需缓存完整 K/V # 从潜变量重建 K/V 的 NoPE 部分推理时从 c_kv 上投影 K_nope c_kv params[W_UK] # [B, T, num_heads * d_k] V c_kv params[W_UV] # [B, T, num_heads * d_v] # Step 2: RoPE 部分位置信息保留 # K 的 RoPE 分量单独从 X 投影不走压缩路径 K_rope RoPE(X params[W_KR]) # [B, T, d_r] 共享不按头区分 # Q 也分两部分低秩压缩 RoPE c_q X params[W_DQ] # Q 低秩压缩 Q_nope c_q params[W_UQ] # [B, T, num_heads * d_k] Q_rope RoPE(c_q params[W_QR]) # [B, T, num_heads * d_r] # Step 3: 拼合完整 Q 和 K # 每个 head 的 Q [Q_nope_h ; Q_rope_h] # 每个 head 的 K [K_nope_h ; K_rope] (K_rope 各 head 共享) Q concat([Q_nope, Q_rope], dim-1) # [B, T, num_heads, d_k d_r] K concat([K_nope, K_rope.expand(num_heads)], dim-1) # Reshape Q Q.view(B, T, num_heads, -1).transpose(1, 2) K K.view(B, T, num_heads, -1).transpose(1, 2) V V.view(B, T, num_heads, d_v).transpose(1, 2) # Step 4: 标准 Attention d_total Q.shape[-1] scores (Q K.transpose(-2, -1)) / sqrt(d_total) attn softmax(scores, dim-1) out attn V # Step 5: 输出 # 优化技巧W_UV 可吸收进 W_O推理时无需显式还原 V out out.transpose(1, 2).view(B, T, -1) return out params[W_O] 关三种架构的 KV Cache 对比AttentionKV cache tokenMHA2 × n_h × d_h (per layer)GQA2 × n_g × d_h (n_g n_h)MLAd_c d_r (d_c n_h × d_h)三、 优化推理效率让 GPU 更忙而不是更闲有了更小的 KV接下来的问题是怎么让 GPU 把这些 KV 用得更高效这里有两个层次的优化计算效率让每一次注意力运算更快 → Flash Attention存储效率让显存的 KV Cache 分配更合理 → Paged Attention两者不是竞争关系而是计算协议和存储协议的配合现代推理框架如 vLLM同时使用两者。3.1 Flash Attention小锅快炒别让 GPU 等数据问题要理解 Flash Attention必须先知道 GPU 有一个经常被忽视的内存层级GPU 的算力极强但数据必须先从 HBM仓库搬到 SRAM厨房才能计算。大量时间不是浪费在算上而是浪费在等搬运上。标准 Attention 有多糟糕它要把 S×S 的注意力矩阵在 HBM 和 SRAM 之间反复搬运① 计算 Q×K^T → 写入 HBMS×S矩阵 ② 读回 HBM → 做 Softmax → 写回 HBM再存一次 ③ 读回 HBM → 乘 V → 写出最终结果解决方案:Tiling Online SoftmaxFlash Attention 的核心思路是小锅快炒每次只从 HBM 取一小块 Q、K、V放进 SRAM 里把打分、Softmax、加权求和全都在 SRAM 里一次性完成只把最终结果写回 HBM。关键数学技巧是 Online SoftmaxSoftmax 本来需要看全所有分数才能归一化但通过维护两个滚动变量当前最大值 m 和累加和 l可以流式地更新归一化结果数学上与完整的 Softmax 完全等价。效果对比指标标准 AttentionFlash AttentionHBM 读写量O(S²)O(S)线性中间矩阵显存占用O(S²)O(1)不需要写出Paged Attention停车场不必都停豪华车位FlashAttention解决的是计算效率PagedAttention解决的是存储管理效率。问题在大模型生成对话时KV Cache键值缓存会随着字数增加而增长。系统不知道你会聊多长所以通常会为每个用户预分配一块连续的大空间。那问题就来了该预先分配多少内存用于Kv Cache呢分配少了不够用分配多了浪费。于是就会出现三种常见的浪费现象预留浪费 我预留了能存 2048 个词的空间结果你只说了 10 个词剩下的 1900 多个位置全空着别人也用不了。碎片浪费 有些空间零零散散虽然加起来很大但因为不连续新来的长请求放不进去。 这导致显存利用率通常只有 20% - 40%非常浪费钱。前缀浪费因为给每个请求都单独分配kv cache空间所以当多个请求有相同前缀时也无法共享而是要重复N份相同的kv Cache存储。解决方案虚拟内存按需分配PagedAttention 借用了计算机系统的“虚拟内存”概念。操作系统物理内存 → 固定大小的页 → 页表映射虚拟地址 ↕ 一一对应 ↕ Paged AttnKV Cache → 固定大小的Block→ Block Table映射逻辑位置具体机制显存切块把全部显存预先切成固定大小的 Block比如每块存 16 个 token 的 KV按需分配一开始只分配 1 个 Block用满再申请下一个不用就立刻释放Block Table维护每个请求的逻辑位置→物理 Block 的映射表物理上可以不连续前缀共享100 个请求的 System Prompt 物理上只存 1 份Block Table 里各自指向同一个物理 BlockCopy-on-Write引入Paged Attention平均预估throughtput能提升8倍以上四、长文本稀疏注意力让需要 Attend 的 Token 变少前面 MLA、GQA 优化的是每个 token 的 KV 有多厚纵向压缩。但面对 1M token 的超长上下文即便每个 token 的 KV 只有 70KB整个 KV Cache 依然高达 70GB更棘手的是计算量标准 Attention 的复杂度是 O(L²)L1M 时就是 10¹² 次运算根本无法承受。所以我们需要另一个维度的压缩横向稀疏化让每个 token 只 attend 最相关的一小部分历史而不是全部。4.1 NSA三镜头摄影系统DeepSeek 稀疏注意力的起点NSA是DeepSeek最早提出的原生可训练的稀疏注意力机制。NSA 的核心思想可以用一个摄影类比来理解 三镜头摄影系统超广角镜压缩分支捕捉整体场景分辨率低但覆盖全长焦镜选择分支锁定远处关键目标精准清晰标准镜滑动窗口拍清近处细节最高分辨率具体的实现方案如下分支一压缩注意力超广角——宏观感知把历史序列按 block_size 分组每组用一个可学习的 MLP 压缩成 1 个向量。Query 和这些摘要向量做注意力快速扫视全局。原始序列: [t1 t2 t3 t4] [t5 t6 t7 t8] ... [t_n-3 t_n-2 t_n-1 t_n] ↓ MLP压缩 ↓ ↓ 摘要向量: [cmp_1] [cmp_2] ... [cmp_k] 计算量: 从 O(L) → O(L/block_size)分支二选择注意力长焦——精准检索压缩分支已经算过每个 block 有多相关了。直接复用这些分数Top-K 选出最相关的块回到原始未压缩的 KV 中取出精细内容再做一次精确注意力。滑动窗口标准镜——局部连贯取最近 win_size 个 token 的完整 KV保证模型对近期上下文的精细理解不受稀疏化影响。最终输出的是以上三路分支拼接后通过门控进行加权的结果让模型学会选择究竟在不同情况下应该使用哪个分支。NSA其实研究性质的论文而后面DeepSeek在V3.2V4两个最新的模型版本中正式把NSA工程化落地。4.2 DSA把 NSA 装进工厂V3.2 的工程验证NSA 是方向DSADeepSeek Sparse Attention是第一次工程落地随 DeepSeek-V3.2 发布。DSA 的核心创新叫 Lightning Indexer——一个用低精度FP8/FP4、低维度远小于 d_model的 Q/K 投影来快速打分的轻量级索引器# 核心Lightning Indexer 实现 class LightningIndexer(nn.Module): def forward(self, x): # 1. 低维投影d_indexer 远小于 d_model q_idx linear_fp8(x, W_qi) k_idx linear_fp8(x, W_ki) # 2. Block 聚合 (工程上通常在投影后做 Pool) k_idx_blocks pool1d(k_idx, kernelblock_size, strideblock_size) # 3. 计算相关性分数 (此处不使用 Softmax使用 ReLU 提高硬件吞吐) scores relu(q_idx k_idx_blocks.T) return scores # 4. Top-K 掩码生成关键工程细节粒度问题DSA 论文层面是 token 粒度打分但由于叠加了 MLA 的 MQA 模式同一 KV entry 被所有 Q head 共享工程 kernel 层面强制变成了 Block 粒度选择才能保证计算效率。ReLU Dotsoftmax的打分机制同样被简化成直接QK乘积过Relu使用FP8或者FP4低精度对齐训练DSA 不是从头训练而是在已有稠密模型上做后训练。为了让 Indexer 学会稠密 Attention 的打分直觉引入了KL 对齐损失让 Lightning Indexer 的输出分布逐渐逼近完整 Attention 的权重分布。DSA更像是一个实验性作品验证NSA提出的稀疏注意力是能跑通的。4.3 V4 的 Hybrid Attention重新排列组合随后DeepSeek V4推出了Hybrid Attention进一步完善了NSA的设计通过把粗粒度记忆、细粒度检索、局部窗口这三种机制重新排列组合。把CSA精准稀疏窗口和HCA粗粒度全局窗口这两种注意力在层与层之间交替使用, 来实现更优的长文本推理效果。CSA先压缩 4×再精准检索DeepSeek V4推出的CSA在DSA的基础上主要增加了三个改进1. 改进一最核心先压缩序列再做稀疏选择DSA: [1M token] → Indexer 打分 → Top-K 选择 CSA: [1M token] → 4× 压缩 → [250K 条目] → Indexer 打分 → Top-K 选择 Indexer 的搜索空间直接缩小 4 倍 KV Cache 的存储量也缩小 4 倍2. 改进二双流重叠压缩消除块切割边界伪影这是 CSA 最精妙的设计。固定分块压缩有一个天然缺陷如果深度学习这个词组恰好被切成 [深度] 和 [学习] 两个 Block后续稀疏选择可能只选了其中一个关键信息就断掉了。CSA 的解法是引入两个位置重叠的KV流流 a 和流 b 不是简单的错位偏移 而是对有重叠的两段隐状态 H用两套独立的投影矩阵产生两套 KV 表示 C_a H W_aKV 位置偏置 B_a ← [0,1,2,3] C_b H W_bKV 位置偏置 B_b ← [2,3,4,5] 关键两个流在 Block 内做跨流联合 Softmax 归一化 模型自动学会这个位置是流 a 表示得更好还是流 b相似的思路你在RAG的overlap-chunking中也能看到。3. 改进三保留滑动窗口分支补回局部精细细节压缩后的远端历史 未压缩的近端 token两者拼接后做最终注意力。HCA压缩 128×换取廉价全局感知HCA 的思路更激进既然 CSA 已经处理精准远端检索那 HCA 就专门负责粗粒度全局感知。也就是NSA压缩压缩注意力的实现。机制压缩比把CSA14的压缩比放大到1128稠密注意力不再使用稀疏注意力而是通过稠密注意力学习粗粒度的全局信息。实现上也没有使用CSA的双流毕竟压缩比例这么大追求边界是否清晰已经没有意义。拼接token粒度滑动窗口注意力层级交替策略三种角色各司其职对于 DeepSeek-V4-Flash前两层使用纯滑动窗口注意力其余层交替使用 CSA 和 HCA对于 DeepSeek-V4-Pro前两层使用 HCA其余层交替使用 CSA 和 HCA。在 1M token 上下文设置下DeepSeek-V4-Pro 仅需 DeepSeek-V3.2 单 token 推理 FLOPs 的 27% 和 10% 的 KV Cache。这一篇就先聊到这里啦~
返回列表