大模型入门：MHA到GQA，一次讲清KV Cache省显存

大模型入门：从 MHA 到 GQA，一文讲清 KV Cache 如何节省显存

很多用户初次运行本地大模型时，往往以为显存的主要消耗来自模型参数。实际上，一个 7B 模型即使使用 FP16 精度，参数显存也需要十几 GB 级别，这当然没错。

然而，进入真实推理场景后，另一个显存消耗因素会快速增长：Prompt 越长，KV Cache 越大；batch 越大，KV Cache 越大；上下文窗口越长，KV Cache 越大；并发请求越多，KV Cache 越难管理。

模型参数在加载时基本固定，而 KV Cache 则随着请求长度和 batch 持续增长。这也正是服务端推理框架重点优化 KV Cache 管理的原因——vLLM 的 PagedAttention、Hugging Face 的 DynamicCache/StaticCache/QuantizedCache，本质上都在解决同一类问题：如何让历史 K/V 既能快速读取，又不会撑爆显存。

GQA 恰好处于这个问题的核心位置。

一句话概括：GQA 通过减少 KV Head 的数量，直接压缩了 KV Cache 的体积。

1. 回顾 KV Cache 到底缓存了什么

Decoder-only 大模型推理通常分为两个阶段：

阶段	输入	主要动作
Prefill	完整 prompt	一次性计算 prompt 的每层 K/V，并写入 cache
Decode	当前新 token	只计算新 token 的 Q/K/V，用新 Q 查询历史 K/V

Hugging Face 的缓存文档也明确指出：自回归生成是一个 token 一个 token 往后预测，KV Cache 会保存过去 token 在注意力层里的 K/V，后续 token 可以复用它们，避免重复计算。

在上一篇文章中，我们使用的 MHA 张量形状如下：

q.shape == [batch, num_heads, seq_len, head_dim]
k.shape == [batch, num_heads, seq_len, head_dim]
v.shape == [batch, num_heads, seq_len, head_dim]

每一层需要缓存的历史 token 的 k 和 v：

past_k.shape == [batch, num_heads, past_len, head_dim]
past_v.shape == [batch, num_heads, past_len, head_dim]

注意，这里缓存的是每一层的 K/V。一个 32 层模型就会有 32 份这样的缓存。因此 KV Cache 的显存可以粗略估算为：

KV Cache bytes = batch_size * seq_len * num_layers * 2 * num_kv_heads * head_dim * bytes_per_element

公式中的 2 代表 K 和 V 两份。其中最容易忽略的变量是 num_kv_heads。

在 MHA 中：num_kv_heads = num_query_heads；而在 GQA 中：num_kv_heads < num_query_heads。这正是 GQA 节省显存的切入点。

2. 用一组数字量化计算

假设一个简化配置：

batch_size = 1
seq_len = 8192
num_layers = 32
num_query_heads = 32
head_dim = 128
dtype = fp16  # 2 bytes

如果是传统 MHA：num_kv_heads = 32，KV Cache 大约为：

1 * 8192 * 32 * 2 * 32 * 128 * 2 bytes = 4 GiB

如果换成 GQA，假设 num_kv_heads = 8，KV Cache 大约为：

1 * 8192 * 32 * 2 * 8 * 128 * 2 bytes = 1 GiB

同样的 Query Head 数量和上下文长度，仅将 KV Head 从 32 降至 8，缓存就缩小为原来的四分之一。

如果是 MQA：num_kv_heads = 1，KV Cache 进一步降至：

128 MiB

这只是一个教学估算，实际框架还会受到 allocator、block size、padding、并发调度、量化和 kernel 实现的影响。但对于面试和工程理解，这个公式足以抓住核心。

3. MHA、MQA、GQA 的区别

可以通过一张表格快速理解：

结构	Query Head	KV Head	直觉
MHA	多个	和 Query 一样多	每个 Q head 独享一组 K/V
MQA	多个	1 个	所有 Q head 共享同一组 K/V
GQA	多个	介于 1 和 Query Head 之间	一组 Q head 共享一组 K/V

假设：num_query_heads = 32，num_kv_heads = 8，则 group_size = num_query_heads // num_kv_heads = 4。

也就是说，前 4 个 Q head（0123）共享一个 KV Head，接下来 4 个（4567）共享下一个，以此类推。它不像 MQA 那样把所有 Query Head 都压到同一个 KV Head 上，也不像 MHA 那样每个 Query Head 都保留独立 K/V。

GQA 原论文的动机也在于此：MQA 可以显著提升 decoder 推理速度，但可能带来质量下降；GQA 使用介于 1 和 Query Head 数之间的 KV Head 数量，在效果和推理效率之间取得折中。

4. 张量形状如何变化

MHA 的投影通常是：

q_proj: hidden_dim -> num_q_heads * head_dim
k_proj: hidden_dim -> num_q_heads * head_dim
v_proj: hidden_dim -> num_q_heads * head_dim

GQA 的投影则变为：

q_proj: hidden_dim -> num_q_heads * head_dim
k_proj: hidden_dim -> num_kv_heads * head_dim
v_proj: hidden_dim -> num_kv_heads * head_dim

可以看到，Q 仍然有很多头，而 K/V 变少了。

假设：batch = 2，seq_len = 5，num_q_heads = 32，num_kv_heads = 8，head_dim = 128，那么：

q.shape == [2, 32, 5, 128]
k.shape == [2, 8, 5, 128]
v.shape == [2, 8, 5, 128]

但 attention 计算时，q @ k.transpose(-2, -1) 要求 head 维度对齐。一个教学做法是将 K/V 按组展开：

k_expanded.shape == [2, 32, 5, 128]
v_expanded.shape == [2, 32, 5, 128]

PyTorch 的 scaled_dot_product_attention(enable_gqa=True) 文档也展示了类似逻辑：启用 GQA 时，会按 Query Head 和 KV Head 的比例对 key/value 做 repeat_interlea ve。但需注意，实际高性能实现不一定物理复制 K/V；服务端推理更关注 cache 布局、访存和 kernel 实现方式。

5. 手写一个最小 GQA 实现

以下代码只保留核心逻辑，适合面试讲解：

• Q Head 数可大于 KV Head 数；
• KV Head 必须能整除 Query Head；
• K/V 先按较少 head 存储；
• 计算 attention 前按组展开；
• cache 中只缓存较少的 KV Head。

import math
import torch
from torch import nn

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    # x: [B, H_kv, T, D]
    if n_rep == 1:
        return x
    batch, num_kv_heads, seq_len, head_dim = x.shape
    x = x[:, :, None, :, :]
    x = x.expand(batch, num_kv_heads, n_rep, seq_len, head_dim)
    return x.reshape(batch, num_kv_heads * n_rep, seq_len, head_dim)

class GroupedQueryAttention(nn.Module):
    def __init__(self,
                 hidden_dim: int,
                 num_q_heads: int,
                 num_kv_heads: int,
                 dropout: float = 0.0,
                 ):
        super().__init__()
        assert hidden_dim % num_q_heads == 0
        assert num_q_heads % num_kv_heads == 0
        self.hidden_dim = hidden_dim
        self.num_q_heads = num_q_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = hidden_dim // num_q_heads
        self.num_groups = num_q_heads // num_kv_heads

        self.q_proj = nn.Linear(hidden_dim, num_q_heads * self.head_dim)
        self.k_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)
        self.v_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)
        self.o_proj = nn.Linear(num_q_heads * self.head_dim, hidden_dim)
        self.dropout = nn.Dropout(dropout)

    def _split_heads(self, x: torch.Tensor, num_heads: int) -> torch.Tensor:
        batch, seq_len, _ = x.shape
        x = x.view(batch, seq_len, num_heads, self.head_dim)
        return x.transpose(1, 2)  # [B, H, T, D]

    def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
        batch, heads, seq_len, head_dim = x.shape
        x = x.transpose(1, 2).contiguous()
        return x.view(batch, seq_len, heads * head_dim)

    def forward(self,
                x: torch.Tensor,
                attn_mask: torch.Tensor | None = None,
                past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
                use_cache: bool = False,
                ):
        q = self._split_heads(self.q_proj(x), self.num_q_heads)
        k = self._split_heads(self.k_proj(x), self.num_kv_heads)
        v = self._split_heads(self.v_proj(x), self.num_kv_heads)

        if past_key_value is not None:
            past_k, past_v = past_key_value
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        present_key_value = (k, v) if use_cache else None

        k_for_attn = repeat_kv(k, self.num_groups)
        v_for_attn = repeat_kv(v, self.num_groups)

        scores = q @ k_for_attn.transpose(-2, -1)
        scores = scores / math.sqrt(self.head_dim)
        if attn_mask is not None:
            scores = scores.masked_fill(attn_mask, float("-inf"))
        weights = torch.softmax(scores, dim=-1)
        weights = self.dropout(weights)

        out = weights @ v_for_attn
        out = self._merge_heads(out)
        out = self.o_proj(out)
        return out, weights, present_key_value

测试一下形状：

x = torch.randn(2, 5, 4096)
gqa = GroupedQueryAttention(
    hidden_dim=4096,
    num_q_heads=32,
    num_kv_heads=8,
)
out, weights, cache = gqa(x, use_cache=True)
print(out.shape)       # [2, 5, 4096]
print(weights.shape)   # [2, 32, 5, 5]
print(cache[0].shape)  # [2, 8, 5, 128]
print(cache[1].shape)  # [2, 8, 5, 128]

关键点在最后两行。注意力权重仍然是 32 个 Query Head：weights.shape == [2, 32, 5, 5]，但缓存中只有 8 个 KV Head：cache[0].shape == [2, 8, 5, 128]，cache[1].shape == [2, 8, 5, 128]。这就是 GQA 在 KV Cache 上节省显存的直接体现。

6. 使用 PyTorch 接口实现 GQA

PyTorch 的 torch.nn.functional.scaled_dot_product_attention 提供了 enable_gqa 参数。

一个最小示例：

import torch
import torch.nn.functional as F

query = torch.randn(2, 32, 5, 128, device="cuda", dtype=torch.float16)
key = torch.randn(2, 8, 5, 128, device="cuda", dtype=torch.float16)
value = torch.randn(2, 8, 5, 128, device="cuda", dtype=torch.float16)

out = F.scaled_dot_product_attention(
    query, key, value,
    is_causal=True,
    enable_gqa=True,
)
print(out.shape)  # [2, 32, 5, 128]

官方文档中有两个重要约束：

number_of_heads_query % number_of_heads_key_value == 0
number_of_heads_key == number_of_heads_value

也就是说：

• Query Head 数必须能被 KV Head 数整除；
• Key Head 数与 Value Head 数必须相同；
• enable_gqa 目前仍是实验特性，后端支持和张量类型有限制。

还有一个容易踩坑的点：PyTorch 该函数中布尔 attn_mask 的语义与某些 MHA 接口的 padding mask 语义相反。scaled_dot_product_attention 中 True 表示参与 attention，迁移代码时需特别留意。

7. 为什么 GQA 主要影响推理

如果只做一次完整 forward 且不使用 KV Cache，GQA 对峰值显存的影响不如 KV Cache 场景那么直观。

真正的收益集中在自回归 decode 阶段：

每一步都要读取历史 K/V
历史越长，读取量越大
并发越高，cache 越多
KV Head 越少，cache 越小

Hugging Face 的优化文档也指出，减少 KV 向量数量只有在使用 KV Cache 的自回归解码场景中才特别有意义，因为 decode 阶段会反复读取历史 K/V，内存带宽很容易成为瓶颈。

因此可以这样理解：

场景	GQA 的价值
训练全序列并行	不是主要优化目标
Prefill	可以减少写入 cache 的 K/V 体积
Decode	最关键，减少每步读取的历史 K/V
长上下文服务	价值更明显
高并发服务	价值更明显

这也就是为什么讲解 GQA 时，不能只关注 attention 公式，更要将其放回推理服务的 KV Cache 场景中审视。

8. 与 vLLM、PagedAttention 的关系

GQA 解决的是单个 token 的 K/V 体积更小；PagedAttention 解决的是大量 token 的 K/V 如何更高效地组织和管理。二者并非同一层优化，但会共同影响推理效率。

vLLM 的 PagedAttention 文档提到，key/value cache 会被拆分为 block，每个 block 存储固定数量 token 的 cache。这样做的目的是用更适合服务端调度的方式管理 KV Cache，而非将每个请求视为一大段连续显存。

可以把它们放在同一张图中：

GQA：减少每个 token 的 KV 体积
PagedAttention：管理大量 token 的 KV 存放方式
Quantized Cache：降低每个元素的字节数
Offloaded Cache：将部分 cache 放到 CPU

如果只看单次模型结构，GQA 是 attention 结构的变化；如果从推理系统角度看，GQA 是 KV Cache 成本控制的重要一环。

9. 常见错误与注意事项

错误 1：只改 num_kv_heads，却忘记修改投影层输出维度

GQA 中 Q/K/V 的 projection 输出维度不同：

q_proj -> num_q_heads * head_dim
k_proj -> num_kv_heads * head_dim
v_proj -> num_kv_heads * head_dim

如果将 K/V 投影到 num_q_heads * head_dim，则 cache 未节省。

错误 2：num_q_heads 不能整除 num_kv_heads

GQA 按组共享 K/V，通常要求 num_q_heads % num_kv_heads == 0，否则每组 Query Head 无法均匀映射到 KV Head。

错误 3：将 repeat 后的 K/V 作为 cache 存储

教学代码为便于理解，在 attention 前做 repeat_kv。但 cache 中应保留较少的 KV Head：cache_k.shape == [B, H_kv, T, D]。若将展开后的 K/V 存入：cache_k.shape == [B, H_q, T, D]，显存又回到 MHA 级别。

错误 4：只计算 cache 容量，忽略内存带宽

KV Cache 不仅占用显存。Decode 每一步都要读取历史 K/V，因此内存带宽也会成为瓶颈。GQA 的价值不仅在于少存，也在于少读。

错误 5：将 GQA 视为无损替换

GQA 是效果与效率的折中。GQA 原论文的结论是，GQA 相比 MQA 更能保留 MHA 的质量，同时接近 MQA 的速度收益。但具体效果仍取决于模型、训练方式、上采样策略和任务。在工程实践中，不应将结构变化理解为“免费优化”——它通常是在模型设计或训练阶段就确定好的。

10. 面试回答要点

如果面试官问：“GQA 和 MHA 有什么区别？”

可以回答：GQA 的核心差异在于，Query Head 的数量多于 Key/Value Head 的数量，多个 Query Head 共享同一组 K/V。而 MHA 中每个 Query Head 都有独立的 K/V。

若继续追问：“为什么能节省显存？”

可以接着回答：因为 KV Cache 的大小与 num_kv_heads 直接成正比。在相同的 Query Head 数量和序列长度下，GQA 只需缓存更少的 K/V Head，因此显存占用更小。

若问：“GQA、MQA 如何区分？”

可以回答：MQA 是所有 Query Head 共享一个 KV Head，极端节省显存但可能损失效果；GQA 是折中方案，将 Query Head 分成若干组，每组共享一个 KV Head。

若问：“代码中最容易出错的地方是什么？”

可以回答：最容易出错的是投影层的输出维度改错，以及 cache 中意外存储了展开后的 K/V。核心约束是 num_q_heads % num_kv_heads == 0。

11. 一张速记表

问题	关键回答
GQA 改了什么？	Query Head 多，KV Head 少
为什么能节省显存？	KV Cache 大小与 num_kv_heads 成正比
MHA 的 KV Head 数？	通常等于 Query Head 数
MQA 的 KV Head 数？	1 个
GQA 的 KV Head 数？	介于 1 和 Query Head 数之间
代码核心约束？	num_q_heads % num_kv_heads == 0
cache 中存储什么？	未展开的 K/V，形状为 [B, H_kv, T, D]
attention 前做什么？	将 K/V 按组映射到 Query Head
最适合讲解的场景？	长上下文、自回归 decode、高并发推理
PyTorch 接口？	scaled_dot_product_attention(..., enable_gqa=True)

总结

GQA 可以用三句话记住：

1. MHA 每个 Query Head 通常都有自己的 K/V，KV Cache 按 Query Head 数增长。
2. GQA 让一组 Query Head 共享较少的 K/V Head，KV Cache 按 KV Head 数增长。
3. 它的主要价值体现在自回归推理，尤其是长上下文和高并发服务场景中。

所以，学习 GQA 不能只记住一个缩写。真正需要记住的是这条逻辑链路：

MHA 张量形状 -> KV Cache 显存公式 -> KV Head 数量 -> Decode 访存压力 -> GQA

这条链路理清之后，GQA、MQA、KV Cache 以及长上下文推理优化就能串联起来。

参考资料

• Joshua Ainslie et al.：GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
  arxiv.org/abs/2305.13…
• PyTorch：torch.nn.functional.scaled_dot_product_attention
  docs.pytorch.org/docs/main/g…
• Hugging Face Transformers：Caching
  huggingface.co/docs/transf…
• Hugging Face Transformers：KV cache strategies
  huggingface.co/docs/transf…
• Hugging Face Transformers：Optimizing LLMs for Speed and Memory
  huggingface.co/docs/transf…
• vLLM：Paged Attention
  docs.vllm.ai/en/latest/d…