Grouped Query Attention（GQA）（グループドクエリアテンション）

Grouped Query Attention（GQA）とは

Grouped Query Attention（GQA）は、2023年にGoogleのAizawa et al.が提案したTransformerのアテンション機構効率化手法です。従来のMulti-Head Attention（MHA）では各ヘッドが独立したQuery・Key・Valueを持ちますが、GQAではQueryヘッドをグループに分け、グループ内で1つのKey/Valueヘッドを共有します。

従来手法との比較

MHA（Multi-Head Attention）

• 全ヘッドが独立したK/V → KVキャッシュが大きい
• 品質は高いが推論コストが高い
• GPT-2・BERT・初期Llamaで採用

MQA（Multi-Query Attention）

• 全QueryヘッドがK/Vを共有 → KVキャッシュ最小
• 速いが品質がMHAより低下する傾向
• Falcon・初期PaLMで採用

GQA（Grouped Query Attention）

• Queryヘッドをグループ分け → グループ内でK/V共有
• MHAとMQAの品質・速度バランスを取る
• Llama 2/3・Mistral・Gemma・Qwenで採用

技術的詳細

グループ構成

標準的な構成例（Llama 3 70B）：

• Queryヘッド数: 64
• Key/Valueヘッド数: 8
• グループ数: 8（各グループ8Qヘッドが1K/Vヘッドを共有）

KVキャッシュ削減効果

MHAと比較してKVキャッシュメモリ使用量を大幅削減できます。 例：32ヘッドMHA → 8グループGQA の場合、KVキャッシュ: 1/4に削減、推論スループット: 最大2〜4倍向上。

実装での採用状況

モデル別GQA設定：Llama 3 8B（Q=32 / KV=8 / グループ=8）、Llama 3 70B（Q=64 / KV=8 / グループ=8）、Mistral 7B（Q=32 / KV=8 / グループ=8）、Gemma 2 9B（Q=16 / KV=8 / グループ=2）、Qwen2 7B（Q=28 / KV=4 / グループ=7）。

HuggingFace Transformersでの実装

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype=torch.float16,
    device_map="auto"
)

config = model.config
print(f"Qヘッド数: {config.num_attention_heads}")
print(f"KVヘッド数: {config.num_key_value_heads}")
print(f"GQAグループ数: {config.num_attention_heads // config.num_key_value_heads}")

vLLMでのGQA最適化

vLLMなどの推論エンジンはGQAを自動検出してPagedAttentionと組み合わせ最適化します：

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Meta-Llama-3-8B-Instruct",
    tensor_parallel_size=1,
    gpu_memory_utilization=0.9
)

sampling_params = SamplingParams(temperature=0.8, max_tokens=512)
outputs = llm.generate(["説明して: "], sampling_params)

GQAの学習とファインチューニング

GQAモデルをLoRAでファインチューニングする際、target_modulesでK/Vを明示的に含めることが重要です：

from peft import LoraConfig

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

GQAはKVキャッシュ削減による長コンテキスト処理の効率化、スループット向上、コスト削減に直結する技術として、LLM推論インフラ設計における必須知識です。現代の主要LLMにおける事実上の標準アテンション機構として、FlashAttention-2との組み合わせで最大の効率を発揮します。