MHA：从古老的Multi Head Attention说起

标准 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.shapegroups 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 对比Attention KV cache tokenMHA 2 × n_h × d_h (per layer)GQA 2 × n_g × d_h (n_g n_h)MLA d_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 有一个经常被忽视的内存层级imageGPU 的算力极强但数据必须先从 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 完全等价。image效果对比指标 标准 Attention Flash 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存储。解决方案虚拟内存按需分配