GQA

import torch
from einops import rearrange, einsum

# 示例维度
batch_size = 1
seq_len = 256
num_query_heads = 8
num_key_value_heads = 2
head_dim = 64

# 初始化随机张量
query = torch.randn(batch_size, seq_len, num_query_heads, head_dim)  # 形状: (1, 256, 8, 64)
key = torch.randn(batch_size, seq_len, num_key_value_heads, head_dim)  # 形状: (1, 256, 2, 64)
value = torch.randn(batch_size, seq_len, num_key_value_heads, head_dim)  # 形状: (1, 256, 2, 64)

# 实际在使用过程中的k和v会被投影到更小的空间

# 计算头组数量
num_head_groups = query.shape[2] // key.shape[2]
print(f"头组数量: {num_head_groups}")

# b: batch_size = 1
# l: seq_len = 256
# q: num_query_heads = 8
# k: num_key_value_heads = 2
# d: head_dim = 64
# g: num_head_groups (num_query_heads/num_key_value_heads) = 4

# 重排张量以高效计算
query = rearrange(query, "b l q d -> b q l d")  # 形状: (1, 8, 256, 64)
key = rearrange(key, "b l k d -> b k l d")  # 形状: (1, 2, 256, 64)
value = rearrange(value, "b l k d -> b k l d")  # 形状: (1, 2, 256, 64)

# 分组查询 (1, 8, 256, 64) -> (1, 4, 2, 256, 64)
query = rearrange(query, "b (k g) l d -> b g k l d", g=num_head_groups)  # 形状: (1, 4, 2, 256, 64)

# 计算注意力分数 (1, 4, 2, 256, 64)@(1, 2, 256, 64) -> (1, 4, 2, 256, 256) -reduce-> (1, 2, 256, 256)
scores = einsum(query, key, "b g k l d, b k l d -> b k l l")  # 形状: (1, 2, 256, 256)
print(f"注意力分数形状: {scores.shape}")

# 最终注意力输出
scale = query.size(-1) ** 0.5
attention = torch.nn.functional.softmax(scores / scale, dim=-1)

# 与值张量相乘 (1, 2, 256, 256)@(1, 2, 256, 64) -> (1, 2, 256, 64)
out = einsum(attention, value, "b k l l, b k l d -> b k l d")  # 形状: (1, 2, 256, 64)

# 重排回原始形状
out = rearrange(out, "b k l d -> b l k d")  # 形状: (1, 256, 2, 64)
print(f"最终输出形状: {out.shape}")