AI演算法注意力機制演進：從MHA到MLA與Flash Attention全解析

AI与大数据

摘要

  • 注意力機制是Transformer的核心：佔LLM 60%+計算量與記憶體，優化注意力＝優化整個模型
  • 5代注意力演進：MHA→MQA→GQA→MLA→Flash Attention，每代都在計算效率與模型品質間尋找平衡
  • MLA（DeepSeek獨創）是2026年最激進的創新：KV Cache壓縮95%+，推理成本降低10倍
  • Flash Attention 3是硬體級優化：H100上實現75%理論峰值利用率，比FA2快2倍
  • 本文提供5代注意力的數學推導、程式碼實作與生產部署對比

目錄

注意力機制：Transformer的心臟

注意力在LLM中的資源佔比

元件 計算量佔比 記憶體佔比 優化價值
Self-Attention 40-50% 50-60% 最高
FFN/MLP 40-50% 25-30% 高
Embedding 2-5% 5-10% 低
LayerNorm <1% <1% 最低

5代注意力演進路線

代際 時間 代表 KV Cache 計算效率 模型品質
第1代 MHA 2017 Transformer 基線 基線 最好
第2代 MQA 2019 PaLM -75% +30% 略降
第2代 GQA 2023 LLaMA-2 -50% +20% 接近MHA
第3代 MLA 2024 DeepSeek-V2 -95% +40% 接近MHA
第4代 FA3 2024 FlashAttn3 不變 +100% 不變
第5代 稀疏 2025-2026 MoBA -80% +50% 中

第1代：MHA多頭注意力

MHA數學公式

` 標準MHA計算流程：

Input X ∈ R^(B×S×D)

Q = XW_Q, K = XW_K, V = XW_V

分頭： Q_i = Q[:, :, i*d_h:(i+1)d_h] for i = 0, ..., h-1 K_i = K[:, :, id_h:(i+1)d_h] V_i = V[:, :, id_h:(i+1)*d_h]

注意力計算： Attn_i = softmax(Q_i K_i^T / √d_h) V_i

合併： Output = Concat(Attn_0, ..., Attn_{h-1}) W_O

KV Cache大小 = 2 × S × h × d_h × sizeof(dtype) = 2 × S × D × sizeof(dtype)

對於70B模型, S=4096, D=8192, FP16: KV Cache = 2 × 4096 × 8192 × 2 = 128MB/層 64層總計 = 8GB `

MHA實作

`python import torch import torch.nn as nn import math

class MultiHeadAttention(nn.Module): def init(self, hidden_size=8192, num_heads=64): super().init() self.num_heads = num_heads self.head_dim = hidden_size // num_heads self.scale = self.head_dim ** -0.5

    self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
    self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
    self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
    self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)

def forward(self, x, attention_mask=None, past_kv=None):
    B, S, D = x.shape
    
    q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
    k = self.k_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
    v = self.v_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
    
    if past_kv is not None:
        past_k, past_v = past_kv
        k = torch.cat([past_k, k], dim=2)
        v = torch.cat([past_v, v], dim=2)
    
    attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scale
    
    if attention_mask is not None:
        attn_weights = attn_weights.masked_fill(attention_mask == 0, float('-inf'))
    
    attn_weights = torch.softmax(attn_weights, dim=-1)
    attn_output = torch.matmul(attn_weights, v)
    
    attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, D)
    return self.o_proj(attn_output), (k, v)

`

MHA的KV Cache開銷

模型 層數 隱層維度 頭數 KV Cache/層 總KV Cache
7B 32 4096 32 2MB 64MB
14B 40 5120 40 2.5MB 100MB
70B 64 8192 64 8MB 512MB
405B 80 16384 128 32MB 2.56GB

第2代：MQA與GQA分組注意力

MQA（Multi-Query Attention）

` MHA vs MQA:

MHA: 每個頭獨立的K和V Q: [h × d_h] K: [h × d_h] V: [h × d_h] KV Cache = 2 × h × d_h × S

MQA: 所有頭共享一組K和V Q: [h × d_h] K: [1 × d_h] V: [1 × d_h] KV Cache = 2 × d_h × S （減少h倍）

例如 h=64: MQA的KV Cache = MHA的 1/64 ≈ 減少98.4% `

GQA（Grouped-Query Attention）

` GQA: g組頭共享K和V

Q: [h × d_h] K: [g × d_h] V: [g × d_h] KV Cache = 2 × g × d_h × S

當 g=1: 退化為MQA 當 g=h: 退化為MHA 當 g=8: KV Cache減少8倍

LLaMA-2 70B使用 g=8: KV Cache = MHA的 8/64 = 1/8 `

GQA實作

`python class GroupedQueryAttention(nn.Module): def init(self, hidden_size=8192, num_heads=64, num_kv_heads=8): super().init() self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = hidden_size // num_heads self.kv_dim = self.num_kv_heads * self.head_dim self.scale = self.head_dim ** -0.5 self.n_rep = num_heads // num_kv_heads

    self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
    self.k_proj = nn.Linear(hidden_size, self.kv_dim, bias=False)
    self.v_proj = nn.Linear(hidden_size, self.kv_dim, bias=False)
    self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)

def _repeat_kv(self, x):
    if self.n_rep == 1:
        return x
    B, g, S, d = x.shape
    return (
        x[:, :, None, :, :]
        .expand(B, g, self.n_rep, S, d)
        .reshape(B, self.num_heads, S, d)
    )

def forward(self, x, attention_mask=None, past_kv=None):
    B, S, D = x.shape
    
    q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
    k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
    v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
    
    if past_kv is not None:
        past_k, past_v = past_kv
        k = torch.cat([past_k, k], dim=2)
        v = torch.cat([past_v, v], dim=2)
    
    k_expanded = self._repeat_kv(k)
    v_expanded = self._repeat_kv(v)
    
    attn_weights = torch.matmul(q, k_expanded.transpose(-2, -1)) * self.scale
    
    if attention_mask is not None:
        attn_weights = attn_weights.masked_fill(attention_mask == 0, float('-inf'))
    
    attn_weights = torch.softmax(attn_weights, dim=-1)
    attn_output = torch.matmul(attn_weights, v_expanded)
    
    attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, D)
    return self.o_proj(attn_output), (k, v)

`

MHA/MQA/GQA對比

方案 KV組數 KV Cache 模型品質 代表模型
MHA h=64 100% 最好 GPT-3
GQA-8 g=8 12.5% 接近MHA LLaMA-2/3
GQA-4 g=4 6.25% 略降 Mistral
MQA g=1 1.56% 降5-8% PaLM

第3代：MLA多頭潛在注意力

MLA核心思想

┌──────────────────────────────────────────────────────────────┐ │ MLA核心創新 │ │ │ │ 傳統MHA/GQA: │ │ K, V直接儲存 → KV Cache大 │ │ │ │ MLA: │ │ 1. 投影到低維潛在空間: c_kv = Compress(X) │ │ c_kv維度 << K,V維度 → Cache壓縮95%+ │ │ │ │ 2. 推理時從潛在空間恢復: │ │ K = W_k_up × c_kv │ │ V = W_v_up × c_kv │ │ │ │ 3. 吸收技術（Absorption）: │ │ Q × K^T = Q × (W_k_up × c_kv)^T │ │ = (Q × W_k_up^T) × c_kv^T │ │ = Q' × c_kv^T │ │ 避免顯式恢復K，直接在低維空間計算注意力 │ └──────────────────────────────────────────────────────────────┘

MLA實作

`python class MultiHeadLatentAttention(nn.Module): def init( self, hidden_size=4096, num_heads=32, kv_latent_dim=512, ): super().init() self.num_heads = num_heads self.head_dim = hidden_size // num_heads self.kv_latent_dim = kv_latent_dim

    self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
    self.kv_compress = nn.Linear(hidden_size, kv_latent_dim, bias=False)
    self.k_up_proj = nn.Linear(kv_latent_dim, hidden_size, bias=False)
    self.v_up_proj = nn.Linear(kv_latent_dim, hidden_size, bias=False)
    self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
    
    self.scale = self.head_dim ** -0.5

def forward(self, x, attention_mask=None, past_c_kv=None):
    B, S, D = x.shape
    
    q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
    
    c_kv = self.kv_compress(x)
    
    if past_c_kv is not None:
        c_kv_full = torch.cat([past_c_kv, c_kv], dim=1)
    else:
        c_kv_full = c_kv
    
    q_absorbed = torch.matmul(
        q.reshape(B * self.num_heads, S, self.head_dim),
        self.k_up_proj.weight.T,
    ).view(B, self.num_heads, S, self.kv_latent_dim)
    
    attn_weights = torch.matmul(
        q_absorbed, c_kv_full.transpose(-2, -1)
    ) * (self.kv_latent_dim ** -0.5)
    
    if attention_mask is not None:
        attn_weights = attn_weights.masked_fill(attention_mask == 0, float('-inf'))
    
    attn_weights = torch.softmax(attn_weights, dim=-1)
    
    attn_to_c = torch.matmul(attn_weights, c_kv_full)
    
    v = self.v_up_proj(attn_to_c.reshape(B, S, self.kv_latent_dim))
    v = v.view(B, S, D)
    
    return self.o_proj(v), c_kv

`

MLA KV Cache對比

方案 KV維度/頭 總KV維度 KV Cache 壓縮比
MHA(h=128) 128 16384 100% 1倍
GQA(g=8) 128 1024 6.25% 16倍
MLA(d_c=512) - 512 3.1% 32倍
MLA(d_c=256) - 256 1.56% 64倍

第4代：Flash Attention硬體級優化

Flash Attention原理

┌──────────────────────────────────────────────────────────────┐ │ Flash Attention核心思想 │ │ │ │ 標準Attention: │ │ Q,K,V → S=QK^T → P=softmax(S) → O=PV │ │ 問題: S和P是S×S矩陣，記憶體O(S²)，且多次讀寫HBM │ │ │ │ Flash Attention: │ │ 1. 分塊計算（Tiling）: 將Q,K,V分成小區塊 │ │ 2. 在SRAM中完成softmax（線上softmax） │ │ 3. 只寫回最終結果O到HBM │ │ │ │ 記憶體: O(S) vs O(S²) │ │ HBM讀寫: O(S²d/N) vs O(S²d) → 減少N倍 │ │ N = SRAM大小 / 區塊大小 │ └──────────────────────────────────────────────────────────────┘

Flash Attention 3特性

特性 FA1 FA2 FA3
GPU支援 A100 A100/H100 H100+
資料型別 FP16/BF16 FP16/BF16 FP16/BF16/FP8
非同步 否 否 是（WGMMA）
管線化 否 否 是
理論利用率 50% 62% 75%
相對速度 1倍 2倍 4倍

Flash Attention使用

`python from flash_attn import flash_attn_func

def flash_attention_forward(q, k, v, causal=True): output = flash_attn_func( q, k, v, causal=causal, softmax_scale=None, ) return output

import torch.nn as nn

class FlashAttentionLayer(nn.Module): def init(self, hidden_size=8192, num_heads=64): super().init() self.num_heads = num_heads self.head_dim = hidden_size // num_heads

    self.qkv_proj = nn.Linear(hidden_size, 3 * hidden_size, bias=False)
    self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)

def forward(self, x):
    B, S, D = x.shape
    
    qkv = self.qkv_proj(x)
    q, k, v = qkv.chunk(3, dim=-1)
    
    q = q.view(B, S, self.num_heads, self.head_dim)
    k = k.view(B, S, self.num_heads, self.head_dim)
    v = v.view(B, S, self.num_heads, self.head_dim)
    
    output = flash_attn_func(q, k, v, causal=True)
    
    output = output.reshape(B, S, D)
    return self.o_proj(output)

`

Flash Attention效能實測

序列長度 標準Attn FA1 FA2 FA3（H100）
1K 2.1ms 0.8ms 0.5ms 0.3ms
4K 28ms 3.2ms 1.8ms 0.9ms
16K OOM 14ms 7ms 3.5ms
32K OOM 32ms 15ms 7ms
128K OOM OOM 68ms 32ms

第5代：稀疏注意力與混合架構

稀疏注意力型別

型別 稀疏模式 計算量 適用場景
區域窗口 固定窗口 O(S×W) 長文件
全域+區域 少量全域token O(S×(W+G)) 通用
膨脹注意力 膨脹卷積模式 O(S×logS) 層次結構
MoBA 混合分塊 O(S×√S) 通用
線性注意力 核方法 O(S×D²) 超長序列

MoBA（Mixture of Block Attention）實作

`python class MoBABlock: def init(self, block_size=256): self.block_size = block_size

def compute_block_importance(self, q, k_blocks):
    scores = []
    for k_block in k_blocks:
        score = torch.matmul(
            q.mean(dim=1),
            k_block.mean(dim=1).T,
        ).max()
        scores.append(score)
    return torch.tensor(scores)

class MoBAAttention(nn.Module): def init(self, hidden_size, num_heads, block_size=256, top_k_blocks=4): super().init() self.num_heads = num_heads self.head_dim = hidden_size // num_heads self.block_size = block_size self.top_k_blocks = top_k_blocks

    self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
    self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
    self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
    self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)

def forward(self, x):
    B, S, D = x.shape
    
    q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
    k = self.k_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
    v = self.v_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
    
    num_blocks = (S + self.block_size - 1) // self.block_size
    k_blocks = k.chunk(num_blocks, dim=2)
    v_blocks = v.chunk(num_blocks, dim=2)
    
    block_importance = self._score_blocks(q, k_blocks)
    top_blocks = torch.topk(block_importance, min(self.top_k_blocks, num_blocks))
    
    selected_k = torch.cat([k_blocks[i] for i in top_blocks.indices], dim=2)
    selected_v = torch.cat([v_blocks[i] for i in top_blocks.indices], dim=2)
    
    attn_weights = torch.matmul(q, selected_k.transpose(-2, -1)) / (self.head_dim ** 0.5)
    attn_weights = torch.softmax(attn_weights, dim=-1)
    output = torch.matmul(attn_weights, selected_v)
    
    output = output.transpose(1, 2).contiguous().view(B, S, D)
    return self.o_proj(output)

def _score_blocks(self, q, k_blocks):
    scores = []
    for k_block in k_blocks:
        score = torch.matmul(q.mean(dim=-1), k_block.mean(dim=-1).transpose(-2, -1)).mean()
        scores.append(score)
    return torch.stack(scores)

`

總結與延伸閱讀

關鍵要點回顧

  1. MHA是基礎：品質最好但KV Cache最大，適合小模型
  2. GQA是平衡點：KV Cache減少8倍，品質接近MHA，LLaMA-2/3標配
  3. MLA最激進：KV Cache壓縮95%+，DeepSeek獨創，推理成本降低10倍
  4. Flash Attention是硬體優化：不改變演算法，4倍加速，與所有方案相容
  5. 稀疏注意力是未來：MoBA等方案適合超長上下文

注意力方案推薦

場景 推薦方案 原因
通用訓練 GQA + Flash Attention 品質與效率平衡
推理服務 MLA + Flash Attention KV Cache最小
超長上下文 MoBA + Flash Attention 計算量可控
小模型 MHA + Flash Attention 品質優先

需要格式化注意力權重的JSON資料？試試我們的JSON格式化工具和cURL轉程式碼。

延伸閱讀

本站提供瀏覽器本地工具，免註冊即可試用 →

#注意力机制#MHA#GQA#MLA#Flash Attention#2026