AIアルゴリズム注意メカニズムの進化：MHAからMLA、Flash Attentionまで完全解説
概要
- 注意メカニズムはTransformerの中核：LLMの計算量とメモリの60%+を占め、注意の最適化＝モデル全体の最適化
- 5世代の注意進化：MHA→MQA→GQA→MLA→Flash Attention、各世代は計算効率とモデル品質のバランスを模索
- MLA（DeepSeek独自）は2026年最も革新的：KV Cache圧縮95%+、推論コスト10分の1削減
- Flash Attention 3はハードウェアレベル最適化：H100で75%の理論ピーク利用率達成、FA2より2倍高速
- 本記事では5世代の注意の数式導出、コード実装、本番デプロイ比較を提供
目次
- 注意メカニズム：Transformerの心臓
- 第1世代：MHAマルチヘッド注意
- 第2世代：MQAとGQAグループ化注意
- 第3世代：MLAマルチヘッド潜在注意
- 第4世代：Flash Attentionハードウェアレベル最適化
- 第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: 全ヘッドが1組のKとVを共有 Q: [h × d_h] K: [1 × d_h] V: [1 × d_h] KV Cache = 2 × d_h × S （h分の1に削減）
例 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分の1に削減
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分の1に削減 │ │ 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) | 長文書 |
| グローバル+局所 | 少数のグローバルトークン | 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)
`
まとめと関連情報
重要ポイントの振り返り
- MHAは基礎：品質は最高だがKV Cacheが最大、小規模モデルに適している
- GQAはバランスポイント：KV Cacheが8分の1に削減、品質はMHAに近く、LLaMA-2/3の標準
- MLAは最も革新的：KV Cache圧縮95%+、DeepSeek独自、推論コスト10分の1削減
- Flash Attentionはハードウェア最適化：アルゴリズムを変更せず、4倍高速化、全方式と互換
- 疎注意は未来：MoBAなどの方式は超長コンテキストに適している
注意方式の推奨
| シナリオ | 推奨方式 | 理由 |
|---|---|---|
| 汎用トレーニング | GQA + Flash Attention | 品質と効率のバランス |
| 推論サービス | MLA + Flash Attention | KV Cache最小 |
| 超長コンテキスト | MoBA + Flash Attention | 計算量制御可能 |
| 小規模モデル | MHA + Flash Attention | 品質優先 |
注意ウェイトのJSONデータをフォーマットする必要がありますか？JSONフォーマッターとcURL to Codeをお試しください。
関連記事
- 大規模モデル推論高速化3エンジン対決 — 推論高速化ソリューション
- 大規模モデル長コンテキスト最適化 — 長コンテキスト注意
- AIチップHBMメモリボトルネック — メモリ最適化
- GQA: Training Generalized Multi-Query Transformer — GQA論文
- DeepSeek-V2: MLA技術レポート — MLA解説
- Flash Attention 3 — FA3論文
ブラウザローカルツールを無料で試す →