为了减少模型的计算了，针对于大模型结构提出来多种设计思路，其中比较有名的有GQA, MLA, MHA,SWA等等， 其中GPT用到的是GQA,所以本文主要讲这个，其他的可以自行了解。
GQA提速原理：GQA 让多个 Query 头共享同一组 Key/Value,把 K/V 投影计算量和 KV Cache 显存从 `num_heads` 组降到 `num_kv_groups` 组(如 12 → 2),是 MHA 和 MQA 之间的折中方案。与 KV Cache 正交,可同时使用;在大模型长序列场景下收益显著,小模型短序列下扩展开销可能抵消计算节省。
# 分组查询注意力(Grouped-Query Attention, GQA)提速原理
> 配套代码:`gpt_with_kv_gqa.py`,核心改动是将 `MultiHeadAttention` 替换为 `GroupedQueryAttention`,通过让多个 Query 头共享同一组 Key/Value 来减少计算量和显存占用。
## 1. 一句话结论
GQA 让 **多个 Query 头共享同一组 Key/Value**,把 K/V 投影的计算量从 `num_heads` 组降到 `num_kv_groups` 组(如 12 头 → 2 组),K/V 缓存显存占用同比减少,是 MHA 和 MQA 之间的折中方案。
---
## 2. 前提:MHA、MQA、GQA 的演进关系
### 2.1 Multi-Head Attention (MHA)
```
num_heads = 12, num_kv_groups = 12
每个 Query 头都有独立的 Key 和 Value:
Q1 → K1, V1
Q2 → K2, V2
...
Q12 → K12, V12
K/V 投影输出维度: num_heads × head_dim = 12 × 64 = 768
K/V 缓存大小: 12 组
```
### 2.2 Multi-Query Attention (MQA)
```
num_heads = 12, num_kv_groups = 1
所有 Query 头共享同一组 Key 和 Value:
Q1, Q2, ..., Q12 → K1, V1 (共享)
K/V 投影输出维度: 1 × head_dim = 64
K/V 缓存大小: 1 组
```
### 2.3 Grouped-Query Attention (GQA)
```
num_heads = 12, num_kv_groups = 2
每 6 个 Query 头共享一组 Key 和 Value:
Q1~Q6 → K1, V1 (共享)
Q7~Q12 → K2, V2 (共享)
K/V 投影输出维度: 2 × head_dim = 128
K/V 缓存大小: 2 组
```
**GQA 是 MHA 和 MQA 的折中**:
- 比 MHA 省计算量和显存(K/V 从 12 组减到 2 组)
- 比 MQA 保留更多表达能力(不同组可以学到不同的 K/V 表示)
---
## 3. 计算量对比
设 `num_heads = 12`, `head_dim = 64`, `num_kv_groups = 2`,序列长度为 T。
### 3.1 K/V 投影层
**MHA** (`gpt_with_kv_cache.py`):
```python
self.W_key = nn.Linear(d_in, d_out) # 768 → 768
self.W_value = nn.Linear(d_in, d_out) # 768 → 768
# 输出维度: 768 = 12 heads × 64
# 参数量: 768 × 768 × 2 = 1,179,648
```
**GQA** (`gpt_with_kv_gqa.py`):
```python
self.W_key = nn.Linear(d_in, num_kv_groups * head_dim) # 768 → 128
self.W_value = nn.Linear(d_in, num_kv_groups * head_dim) # 768 → 128
# 输出维度: 128 = 2 groups × 64
# 参数量: 768 × 128 × 2 = 196,608
```
**节省**:
- K/V 投影参数量: 1,179,648 → 196,608,**减少 83.3%**
- K/V 投影 FLOPs: 同比例减少
### 3.2 KV Cache 显存
每层每个 token 需要缓存的 K/V 大小:
**MHA**:
```
K: num_heads × head_dim = 12 × 64 = 768 floats
V: num_heads × head_dim = 12 × 64 = 768 floats
总计: 1536 floats = 3072 bytes (FP16)
```
**GQA** (`num_kv_groups = 2`):
```
K: num_kv_groups × head_dim = 2 × 64 = 128 floats
V: num_kv_groups × head_dim = 2 × 64 = 128 floats
总计: 256 floats = 512 bytes (FP16)
```
**节省**:
- KV Cache 显存: 3072 → 512 bytes/token/layer,**减少 83.3%**
- 对于 12 层模型,200 tokens 序列:
- MHA: 3072 × 200 × 12 ≈ **7.4 MB**
- GQA: 512 × 200 × 12 ≈ **1.2 MB**
### 3.3 注意力计算
```python
# MHA
queries: (b, 12, T, 64)
keys: (b, 12, T, 64)
attn_scores = queries @ keys.transpose(2, 3) # (b, 12, T, T)
# GQA
queries: (b, 12, T, 64)
keys: (b, 12, T, 64) # 通过 repeat_interleave 扩展到 12 头
attn_scores = queries @ keys.transpose(2, 3) # (b, 12, T, T)
```
**注意力计算本身完全一样**(都是 12 头 × T × T 的点积),GQA 只是通过共享 K/V 减少了 K/V 投影和缓存的开销。
---
## 4. 代码实现拆解
### 4.1 关键参数
```python
class GroupedQueryAttention(nn.Module):
def __init__(self, d_in, d_out, dropout, num_heads, num_kv_groups, ...):
...
self.num_heads = num_heads # 12
self.head_dim = d_out // num_heads # 64
self.num_kv_groups = num_kv_groups # 2
self.group_size = num_heads // num_kv_groups # 6 (每组几个 query)
```
### 4.2 K/V 投影维度缩减
```python
# MHA (gpt_with_kv_cache.py 第 32-33 行)
self.W_key = nn.Linear(d_in, d_out) # 768 → 768
self.W_value = nn.Linear(d_in, d_out) # 768 → 768
# GQA (gpt_with_kv_gqa.py 第 32-33 行)
self.W_key = nn.Linear(d_in, num_kv_groups * self.head_dim) # 768 → 128
self.W_value = nn.Linear(d_in, num_kv_groups * self.head_dim) # 768 → 128
```
**这就是省计算量的地方**:K/V 投影的输出维度从 `num_heads × head_dim` 缩减到 `num_kv_groups × head_dim`。
### 4.3 分组与扩展
```python
# forward 第 54-56 行:reshape 成分组结构
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
# (b, T, 12, 64) → (b, 12, T, 64)
keys_new = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
# (b, T, 2, 64) → (b, 2, T, 64)
values_new = values.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
# (b, T, 2, 64) → (b, 2, T, 64)
```
```python
# forward 第 73-74 行:通过 repeat_interleave 扩展到与 query 头数匹配
keys = keys_base.repeat_interleave(self.group_size, dim=1)
# (b, 2, T, 64) → (b, 12, T, 64)
# [K1, K2] → [K1, K1, K1, K1, K1, K1, K2, K2, K2, K2, K2, K2]
values = values_base.repeat_interleave(self.group_size, dim=1)
# (b, 2, T, 64) → (b, 12, T, 64)
```
**为什么用 `repeat_interleave` 而不是 `repeat`?**
```python
# repeat_interleave(dim=1): 沿维度 1 重复每个元素
# [K1, K2] → [K1, K1, K1, K1, K1, K1, K2, K2, K2, K2, K2, K2]
# 保证 Q1~Q6 对应 K1,Q7~Q12 对应 K2
# repeat(1, 6, 1, 1): 整体重复 6 次
# [K1, K2] → [K1, K2, K1, K2, K1, K2, K1, K2, K1, K2, K1, K2]
# 这样 Q1 对应 K1,Q2 对应 K2,... 分组关系就乱了
```
### 4.4 KV Cache 拼接维度
```python
# MHA (gpt_with_kv_cache.py 第 60-61 行)
self.cache_k = torch.cat([self.cache_k, keys_new], dim=1)
# dim=1 是 sequence 维度: (b, seq, heads, dim)
# GQA (gpt_with_kv_gqa.py 第 62-63 行)
self.cache_k = torch.cat([self.cache_k, keys_new], dim=2)
# dim=2 是 sequence 维度: (b, heads, seq, dim)
```
**注意**:GQA 的 cache 维度是 `(b, num_kv_groups, seq, dim)` 而不是 `(b, num_heads, seq, dim)`,因为 cache 只存 2 组 K/V,不存 12 组。扩展操作 `repeat_interleave` 在拼接 cache 之后执行。
---
## 5. 与 KV Cache 的关系
GQA 和 KV Cache 是**正交的两个优化**:
| 优化维度 | 解决什么问题 | 节省什么 |
|---|---|---|
| **KV Cache** | 避免重复计算历史 token 的 K/V | 计算量(O(N²) → O(N)) |
| **GQA** | 减少 K/V 投影和缓存的维度 | 参数量 + 显存 |
`gpt_with_kv_gqa.py` **同时使用了两种优化**:
- 有 KV Cache(每步只算 1 个新 token)
- 有 GQA(K/V 从 12 组减到 2 组)
---
## 6. 实际性能对比
以你的 benchmark 结果为例(prompt=32, 生成 200 tokens):
| 模型 | FLOPs | 吞吐量 | 显存峰值 |
|---|---|---|---|
| Vanilla (MHA, no cache) | 6.65e+12 | 102.1 tok/s | 0.75 GB |
| KV-Cache (MHA + cache) | 5.83e+10 | 104.5 tok/s | 0.68 GB |
| **GQA + cache** | **5.29e+10** | 97.0 tok/s | **0.68 GB** |
**FLOPs 对比**:
- GQA vs MHA(都带 cache): 5.29e+10 vs 5.83e+10,**节省 9.3%**
- 这个 9.3% 主要来自 K/V 投影的参数量减少
**显存对比**:
- GQA 和 MHA 的显存峰值相同(0.68 GB),因为对于 124M 小模型,KV Cache 占比很小
- 在 70B 大模型上,KV Cache 显存占比很大,GQA 的显存节省会非常明显
**吞吐量对比**:
- GQA (97.0 tok/s) 略低于 MHA (104.5 tok/s)
- 原因:`repeat_interleave` 扩展操作引入额外开销,在小模型上抵消了计算量节省
- 在大模型 + 长序列场景下,计算量节省会超过扩展开销,GQA 会更快
---
## 7. 与 MHA 的表达能力差异
### 7.1 MHA
每个 Query 头有独立的 K/V,可以学到完全不同的注意力模式:
- Head 1: 关注局部语法关系
- Head 2: 关注长距离指代
- Head 3: 关注特定实体类型
- ...
### 7.2 GQA
多个 Query 头共享 K/V,表达能力略有损失:
- Q1~Q6 共享 K1/V1:这 6 个头的注意力模式会更相似
- Q7~Q12 共享 K2/V2:这 6 个头的注意力模式会更相似
### 7.3 实践经验
- **小模型(< 1B)**:GQA 的表达能力损失可能明显
- **大模型(≥ 7B)**:模型容量充足,GQA 损失可忽略,甚至可能起到正则化效果
- **num_kv_groups 的选择**:
- 2~4 组:平衡计算效率和表达能力(Llama 2 70B 用 8 组)
- 1 组:即 MQA,最激进但损失最大
---
## 8. 一句话总结
> GQA 让多个 Query 头共享同一组 Key/Value,把 K/V 投影计算量和 KV Cache 显存从 `num_heads` 组降到 `num_kv_groups` 组(如 12 → 2),是 MHA 和 MQA 之间的折中方案。与 KV Cache 正交,可同时使用;在大模型长序列场景下收益显著,小模型短序列下扩展开销可能抵消计算节省。代码实现如下：
python
import argparse
import time
import tiktoken
import torch
import torch.nn as nn
#####################################
# NEW: GQA instead of MHA
#####################################
class GroupedQueryAttention(nn.Module):
def __init__(
self, d_in, d_out, dropout, num_heads, num_kv_groups, dtype=None, qkv_bias=False
):
super().__init__()
assert d_out % num_heads == 0, "d_out must be divisible by num_heads"
assert num_heads % num_kv_groups == 0, "num_heads must be divisible by num_kv_groups"
self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads
self.W_key = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=qkv_bias, dtype=dtype)
self.W_value = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=qkv_bias, dtype=dtype)
self.num_kv_groups = num_kv_groups
self.group_size = num_heads // num_kv_groups
self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias, dtype=dtype)
self.out_proj = nn.Linear(d_out, d_out, bias=False, dtype=dtype)
self.dropout = nn.Dropout(dropout)
self.register_buffer("cache_k", None, persistent=False)
self.register_buffer("cache_v", None, persistent=False)
self.ptr_current_pos = 0
def forward(self, x, use_cache=False):
b, num_tokens, _ = x.shape
# Apply projections
queries = self.W_query(x) # (b, num_tokens, num_heads * head_dim)
keys = self.W_key(x) # (b, num_tokens, num_kv_groups * head_dim)
values = self.W_value(x) # (b, num_tokens, num_kv_groups * head_dim)
# Reshape
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
keys_new = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
values_new = values.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
if use_cache:
if self.cache_k is None:
self.cache_k, self.cache_v = keys_new, values_new
else:
self.cache_k = torch.cat([self.cache_k, keys_new], dim=2)
self.cache_v = torch.cat([self.cache_v, values_new], dim=2)
keys_base, values_base = self.cache_k, self.cache_v
else:
keys_base, values_base = keys_new, values_new
if self.cache_k is not None or self.cache_v is not None:
self.cache_k, self.cache_v = None, None
self.ptr_current_pos = 0
# Expand keys and values to match the number of heads
# Shape: (b, num_heads, num_tokens, head_dim)
keys = keys_base.repeat_interleave(self.group_size, dim=1) # Shape: (b, num_heads, num_tokens, head_dim)
values = values_base.repeat_interleave(self.group_size, dim=1) # Shape: (b, num_heads, num_tokens, head_dim)
# For example, before repeat_interleave along dim=1 (query groups):
# [K1, K2]
# After repeat_interleave (each query group is repeated group_size times):
# [K1, K1, K2, K2]
# If we used regular repeat instead of repeat_interleave, we'd get:
# [K1, K2, K1, K2]
# Compute scaled dot-product attention (aka self-attention) with a causal mask
# Shape: (b, num_heads, num_tokens, num_tokens)
attn_scores = queries @ keys.transpose(2, 3) # Dot product for each head
####################################################
# causal mask
num_tokens_Q = queries.shape[-2]
num_tokens_K = keys.shape[-2]
device = queries.device
if use_cache:
q_positions = torch.arange(
self.ptr_current_pos,
self.ptr_current_pos + num_tokens_Q,
device=device,
dtype=torch.long,
)
self.ptr_current_pos += num_tokens_Q
else:
q_positions = torch.arange(num_tokens_Q, device=device, dtype=torch.long)
self.ptr_current_pos = 0
k_positions = torch.arange(num_tokens_K, device=device, dtype=torch.long)
mask = q_positions.unsqueeze(-1) < k_positions.unsqueeze(0)
# Use the mask to fill attention scores
attn_scores = attn_scores.masked_fill(mask, -torch.inf)
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
assert keys.shape[-1] == self.head_dim
attn_weights = self.dropout(attn_weights)
# Shape: (b, num_tokens, num_heads, head_dim)
context_vec = (attn_weights @ values).transpose(1, 2)
# Combine heads, where self.d_out = self.num_heads * self.head_dim
context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
context_vec = self.out_proj(context_vec) # optional projection
return context_vec
def reset_cache(self):
self.cache_k, self.cache_v = None, None
self.ptr_current_pos = 0
#####################################
# Chapter 4
#####################################
class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))
def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift
class GELU(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return 0.5 * x * (1 + torch.tanh(
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
(x + 0.044715 * torch.pow(x, 3))
))
class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)
def forward(self, x):
return self.layers(x)
class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = GroupedQueryAttention(
d_in=cfg["emb_dim"],
d_out=cfg["emb_dim"],
num_heads=cfg["n_heads"],
num_kv_groups=cfg["n_kv_groups"],
dropout=cfg["drop_rate"],
qkv_bias=cfg["qkv_bias"])
self.ff = FeedForward(cfg)
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])
def forward(self, x, use_cache=False):
# Shortcut connection for attention block
shortcut = x
x = self.norm1(x)
# x = self.att(x) # Shape [batch_size, num_tokens, emb_size]
####################################################
# KV cache-related
x = self.att(x, use_cache=use_cache)
####################################################
x = self.drop_shortcut(x)
x = x + shortcut # Add the original input back
# Shortcut connection for feed-forward block
shortcut = x
x = self.norm2(x)
x = self.ff(x)
x = self.drop_shortcut(x)
x = x + shortcut # Add the original input back
return x
class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
self.drop_emb = nn.Dropout(cfg["drop_rate"])
# self.trf_blocks = nn.Sequential(
# *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])
####################################################
# KV cache-related
self.trf_blocks = nn.ModuleList(
[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])
self.current_pos = 0
####################################################
self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
def forward(self, in_idx, use_cache=False):
batch_size, seq_len = in_idx.shape
tok_embeds = self.tok_emb(in_idx)
# pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
####################################################
# KV cache-related
if use_cache:
pos_ids = torch.arange(self.current_pos, self.current_pos + seq_len, device=in_idx.device, dtype=torch.long)
self.current_pos += seq_len
else:
pos_ids = torch.arange(0, seq_len, device=in_idx.device, dtype=torch.long)
pos_embeds = self.pos_emb(pos_ids).unsqueeze(0)
####################################################
x = tok_embeds + pos_embeds # Shape [batch_size, num_tokens, emb_size]
x = self.drop_emb(x)
# x = self.trf_blocks(x)
####################################################
# KV cache-related
for blk in self.trf_blocks:
x = blk(x, use_cache=use_cache)
####################################################
x = self.final_norm(x)
logits = self.out_head(x)
return logits
####################################################
# KV cache-related
def reset_kv_cache(self):
for blk in self.trf_blocks:
blk.att.reset_cache()
self.current_pos = 0
####################################################
def generate_text_simple_cached(model, idx, max_new_tokens,
context_size=None, use_cache=True):
model.eval()
ctx_len = context_size or model.pos_emb.num_embeddings
with torch.no_grad():
if use_cache:
# Init cache with full prompt
model.reset_kv_cache()
logits = model(idx[:, -ctx_len:], use_cache=True)
for _ in range(max_new_tokens):
# a) pick the token with the highest log-probability (greedy sampling)
next_idx = logits[:, -1].argmax(dim=-1, keepdim=True)
# b) append it to the running sequence
idx = torch.cat([idx, next_idx], dim=1)
# c) feed model only the new token
logits = model(next_idx, use_cache=True)
else:
for _ in range(max_new_tokens):
logits = model(idx[:, -ctx_len:], use_cache=False)
next_idx = logits[:, -1].argmax(dim=-1, keepdim=True)
idx = torch.cat([idx, next_idx], dim=1)
return idx
def main():
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, description="Run GPT with grouped-query attention.")
parser.add_argument("--emb_dim", type=int, default=768, help="Model embedding dimension.")
parser.add_argument("--n_heads", type=int, default=12, help="Number of attention heads.")
parser.add_argument("--n_layers", type=int, default=12, help="Number of transformer blocks.")
parser.add_argument("--n_kv_groups", type=int, default=2, help="Number of key/value groups.")
parser.add_argument("--max_new_tokens", type=int, default=200, help="Number of tokens to generate.")
args = parser.parse_args()
start_context = "Hello, I am"
tokenizer = tiktoken.get_encoding("gpt2")
encoded = tokenizer.encode(start_context)
GPT_CONFIG_124M = {
"vocab_size": 50257, # Vocabulary size
"context_length": args.max_new_tokens + len(encoded),
"emb_dim": args.emb_dim, # Embedding dimension
"n_heads": args.n_heads, # Number of attention heads
"n_layers": args.n_layers, # Number of layers
"drop_rate": 0.0, # Dropout rate
"qkv_bias": False, # Query-Key-Value bias
"n_kv_groups": args.n_kv_groups
}
torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device, dtype=torch.bfloat16)
model.eval() # disable dropout
encoded_tensor = torch.tensor(encoded, device=device).unsqueeze(0)
print(f"\n{50*'='}\n{22*' '}IN\n{50*'='}")
print("\nInput text:", start_context)
print("Encoded input text:", encoded)
print("encoded_tensor.shape:", encoded_tensor.shape)
if torch.cuda.is_available():
torch.cuda.synchronize()
start = time.time()
token_ids = generate_text_simple_cached(
model=model,
idx=encoded_tensor,
max_new_tokens=args.max_new_tokens,
)
if torch.cuda.is_available():
torch.cuda.synchronize()
total_time = time.time() - start
decoded_text = tokenizer.decode(token_ids.squeeze(0).tolist())
print(f"\n\n{50*'='}\n{22*' '}OUT\n{50*'='}")
print("\nOutput:", token_ids)
print("Output length:", len(token_ids[0]))
print("Output text:", decoded_text)
print(f"\nTime: {total_time:.2f} sec")
print(f"{int(len(token_ids[0])/total_time)} tokens/sec")
if torch.cuda.is_available():
max_mem_bytes = torch.cuda.max_memory_allocated()
max_mem_gb = max_mem_bytes / (1024 ** 3)
print(f"Max memory allocated: {max_mem_gb:.2f} GB")
if __name__ == "__main__":
main()