実装で理解するMHA・GQA・MQAの違い

Last updated at 2026-07-11Posted at 2026-07-11

はじめに

個人学習・備忘録として書いている記事です。
前回までの記事で、基本的な Attention 機構である MHA (Multi-Head Attention), MQA (Multi-Query Attention), GQA (Grouped-Query Attention) について見てみました。

方式	Query Head数	KV Head数	特徴
MHA	H	H	各Query Headが個別のK、Vを持つ
GQA	H	G	複数のQuery HeadがK、Vを共有する
MQA	H	1	すべてのQuery Headが1組のK、Vを共有する

本稿では、NumPy によるシンプルな実装を通して、それぞれの Attention の shape の変化を追いながら見てみたいと思います。

Multi-Head Attention (MHA)

典型的な MHA のクラス

import numpy as np
class MultiHeadAttention:
    def __init__(self, hidden_size, num_attention_heads):
        assert hidden_size % num_attention_heads == 0

        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.head_dim = hidden_size // num_attention_heads

        self.Wq = np.random.randn(hidden_size, hidden_size) * 0.01
        self.Wk = np.random.randn(hidden_size, hidden_size) * 0.01
        self.Wv = np.random.randn(hidden_size, hidden_size) * 0.01
        self.Wo = np.random.randn(hidden_size, hidden_size) * 0.01

    def split_heads(self, x):
        # x: (B, T, hidden_size)
        B, T, _ = x.shape
        x = x.reshape(B, T, self.num_attention_heads, self.head_dim)
        return x.transpose(0, 2, 1, 3)
        # (B, H, T, head_dim)

    def combine_heads(self, x):
        # x: (B, H, T, head_dim)
        B, H, T, D = x.shape
        x = x.transpose(0, 2, 1, 3)
        return x.reshape(B, T, H * D)
        # (B, T, hidden_size)

    def forward(self, x):
        # x: (B, T, hidden_size)
        print("========== Multi-Head Attention ==========")
        print(f"Input x          : {x.shape}")
 
        Q = x @ self.Wq
        K = x @ self.Wk
        V = x @ self.Wv
        print("\nAfter linear projection")
        print(f"Q                : {Q.shape}")
        print(f"K                : {K.shape}")
        print(f"V                : {V.shape}")

        Q = self.split_heads(Q)
        K = self.split_heads(K)
        V = self.split_heads(V)
        print("\nAfter split_heads()")
        print(f"Q                : {Q.shape}")
        print(f"K                : {K.shape}")
        print(f"V                : {V.shape}")

        scores = Q @ K.transpose(0, 1, 3, 2)
        scores = scores / np.sqrt(self.head_dim)
        print("\nAttention score")
        print(f"K^T              : {K.transpose(0,1,3,2).shape}")
        print(f"Scores           : {scores.shape}")

        attn_weights = softmax(scores, axis=-1)
        print("\nSoftmax")
        print(f"Attention Weight : {attn_weights.shape}")

        context = attn_weights @ V
        print("\nContext")
        print(f"Context          : {context.shape}")

        context = self.combine_heads(context)
        print("\nAfter combine_heads()")
        print(f"Context          : {context.shape}")
        
        output = context @ self.Wo
        print("\nOutput")
        print(f"Output           : {output.shape}")
        print("==========================================\n")

        return output, attn_weights

MHA では、線形変換後の Q・K・V はすべて (B, T, hidden_size) になります。その後、各 Head に分割すると、それぞれ (B, H, T, head_dim) となります。

具体的には、

B = 1
T = 16
hidden_size = 32
head_dim = 4

とすると

(B, T, hidden_size) → (B, num_heads, T, head_dim)

となるため、

(B, T, 32) → (B, 8, T, 4)

となります。(num_heads は hidden_size / head_dim = 8)

以下の条件で試すと次元数を確認できます。

B = 1
T = 16
hidden_size = 32
num_attention_heads = 8
num_kv_heads = 2

x = np.random.randn(B, T, hidden_size)

mha = MultiHeadAttention(hidden_size, num_attention_heads)
output_mha, attn_weights_mha = mha.forward(x)

========== Multi-Head Attention ==========
Input x          : (1, 16, 32)

After linear projection
Q                : (1, 16, 32)
K                : (1, 16, 32)
V                : (1, 16, 32)

After split_heads()
Q                : (1, 8, 16, 4)
K                : (1, 8, 16, 4)
V                : (1, 8, 16, 4)

Attention score
K^T              : (1, 8, 4, 16)
Scores           : (1, 8, 16, 16)

Softmax
Attention Weight : (1, 8, 16, 16)

Context
Context          : (1, 8, 16, 4)

After combine_heads()
Context          : (1, 16, 32)

Output
Output           : (1, 16, 32)
==========================================

Grouped-Query Attention (GQA)

GQA では、K, V のみ Head 数が変わります。
そのため、 MHA における split_heads を拡張して、split_query_heads と split_kv_heads に分けた実装とします。

class GroupedQueryAttention:
    def __init__(
        self,
        hidden_size,
        num_query_heads,
        num_kv_heads,
    ):
        assert hidden_size % num_query_heads == 0
        assert num_query_heads % num_kv_heads == 0
        self.hidden_size = hidden_size
        self.num_query_heads = num_query_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = hidden_size // num_query_heads

        self.Wq = np.random.randn(
            hidden_size,
            num_query_heads * self.head_dim,
        ) * 0.01
        self.Wk = np.random.randn(
            hidden_size,
            num_kv_heads * self.head_dim,
        ) * 0.01
        self.Wv = np.random.randn(
            hidden_size,
            num_kv_heads * self.head_dim,
        ) * 0.01
        self.Wo = np.random.randn(
            hidden_size,
            hidden_size
        ) * 0.01


    def split_query_heads(self, x):
        B, T, _ = x.shape
        x = x.reshape(
            B, T, self.num_query_heads, self.head_dim,
        )
        return x.transpose(0, 2, 1, 3)
    
    def split_kv_heads(self, x):
        B, T, _ = x.shape
        x = x.reshape(
            B, T, self.num_kv_heads, self.head_dim,
        )
        return x.transpose(0, 2, 1, 3)

    def combine_heads(self, x):
        # x: (B, H, T, head_dim)
        B, H, T, D = x.shape
        x = x.transpose(0, 2, 1, 3)
        return x.reshape(B, T, H * D)
        # (B, T, hidden_size)

    def forward(self, x):
        print("========== Grouped-Query Attention ==========")
        print(f"Input x          : {x.shape}")
 
        Q = x @ self.Wq
        K = x @ self.Wk
        V = x @ self.Wv
        print("\nAfter linear projection")
        print(f"Q                : {Q.shape}")
        print(f"K                : {K.shape}")
        print(f"V                : {V.shape}")

        Q = self.split_query_heads(Q)
        K = self.split_kv_heads(K)
        V = self.split_kv_heads(V)
        print("\nAfter split_heads()")
        print(f"Q                : {Q.shape}")
        print(f"K                : {K.shape}")
        print(f"V                : {V.shape}")

        repeat = self.num_query_heads // self.num_kv_heads

        K = np.repeat(K, repeat, axis=1)
        V = np.repeat(V, repeat, axis=1)

        scores = Q @ K.transpose(0, 1, 3, 2)
        scores = scores / np.sqrt(self.head_dim)
        print("\nAttention score")
        print(f"K^T              : {K.transpose(0,1,3,2).shape}")
        print(f"Scores           : {scores.shape}")

        attn_weights = softmax(scores, axis=-1)
        print("\nSoftmax")
        print(f"Attention Weight : {attn_weights.shape}")

        context = attn_weights @ V
        print("\nContext")
        print(f"Context          : {context.shape}")

        context = self.combine_heads(context)
        print("\nAfter combine_heads()")
        print(f"Context          : {context.shape}")

        output = context @ self.Wo
        print("\nOutput")
        print(f"Output           : {output.shape}")
        print("==========================================\n")

        return output, attn_weights

K, V に関しては

hidden_size = 32
num_query_heads = 8
head_dim = 4
num_kv_heads = 2

num_kv_heads × head_dim = 2 x 4 = 8

となるため、

(B, T, 32)
->
(B, 2, T, 4)

となります。
ただし、このままではもとの入力に対する Head 数が足りないので、np.repeat ¹ をつかって Query の次元 (B, 8, T, 4) に対応付けます。
※この実装では説明を分かりやすくするため、np.repeat を使って K, V を Query Head 数まで複製しています。実際の LLM 実装では複製するのではなく、同じ K, V を複数の Query Head が参照するように実装されることが一般的なようです。

先程と同様の入力を使って次元を確かめると以下のようになります。

B = 1
T = 16
hidden_size = 32
num_attention_heads = 8
num_kv_heads = 2

x = np.random.randn(B, T, hidden_size)

gqa = GroupedQueryAttention(hidden_size, num_attention_heads, num_kv_heads)
output_gqa, attn_weights_gqa = gqa.forward(x)

========== Grouped-Query Attention ==========
Input x          : (1, 16, 32)

After linear projection
Q                : (1, 16, 32)
K                : (1, 16, 8)
V                : (1, 16, 8)

After split_heads()
Q                : (1, 8, 16, 4)
K                : (1, 2, 16, 4)
V                : (1, 2, 16, 4)

Attention score
K^T              : (1, 8, 4, 16)
Scores           : (1, 8, 16, 16)

Softmax
Attention Weight : (1, 8, 16, 16)

Context
Context          : (1, 8, 16, 4)

After combine_heads()
Context          : (1, 16, 32)

Output
Output           : (1, 16, 32)
==========================================

先程の MHA と比較して、各 Head における K, V の次元が削減されていることが分かります。

Multi-Query Attention (MQA)

MQA は GQA の特殊ケースであり、num_kv_heads = 1 としたものです。そのため、K, V は 1 組だけ生成され、すべての Query Head がそれらを共有します。
そのため、結果は以下のようになり、K (1,16,4) であることが分かります。

========== Multi-Query Attention ==========
Input x          : (1, 16, 32)

After linear projection
Q                : (1, 16, 32)
K                : (1, 16, 4)
V                : (1, 16, 4)

After split_heads()
Q                : (1, 8, 16, 4)
K                : (1, 1, 16, 4)
V                : (1, 1, 16, 4)

Attention score
K^T              : (1, 8, 4, 16)
Scores           : (1, 8, 16, 16)

Softmax
Attention Weight : (1, 8, 16, 16)

Context
Context          : (1, 8, 16, 4)

After combine_heads()
Context          : (1, 16, 32)

Output
Output           : (1, 16, 32)

まとめ

今回の実装では、MHA・GQA・MQA の違いは K, V の Head 数であることが確認できました。

方式	Q	K/V
MHA	H	H
GQA	H	G
MQA	H	1

その結果、

MHA は最も表現力が高い
GQA は品質とメモリ使用量のバランスが良い
MQA は KV Cache を最も小さくできる

という特徴があります。

https://numpy.org/doc/stable/reference/generated/numpy.repeat.html ↩

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up