大模型推理加速核心：从 MHA 到 GQA 再到 MQA，KV Cache 优化机制全解析

为什么 KV Cache 成了大模型推理的命门？

在 Transformer 架构的大语言模型中，自回归生成（autoregressive generation）是推理阶段的核心模式——每生成一个新 token，模型都需要关注之前所有已生成的 token。这意味着，随着序列长度增长，注意力计算所需的 Key 和 Value 矩阵会线性膨胀，而存储这些矩阵的 KV Cache 就成了显存占用的绝对主力。

以 LLaMA-2-70B 为例，在 FP16 精度下，一个长度为 4096 的序列，KV Cache 占用的显存高达约 20GB，几乎与模型权重本身的显存占用相当。在批量推理场景下，KV Cache 甚至可以轻松吃掉 80% 以上的可用显存。这直接导致两个严重后果：一是单卡能承载的并发请求数极少，二是长上下文推理几乎无法实现。

因此，KV Cache 优化不仅仅是”省点显存”的锦上添花，而是大模型推理从实验室走向生产的关键工程瓶颈。本文将深入剖析从标准多头注意力（MHA）到分组查询注意力（GQA）再到多查询注意力（MQA）的完整演进路径，揭示每种方案背后的设计权衡与工程实践。

深度学习推理优化示意图

MHA 标准架构：KV Cache 膨胀的根源

标准的多头注意力（Multi-Head Attention, MHA）是 Transformer 的原始设计。在这种架构中，每个注意力头都拥有独立的 Query、Key 和 Value 投影。对于模型配置中的

1	num_attention_heads

个头，每个头都会产生一组 Key 和 Value 向量，需要完整存储在 KV Cache 中。

具体来说，假设模型有以下配置：

隐藏维度
1
d_model = 4096
注意力头数
1
n_heads = 32
每头维度
1
d_head = d_model / n_heads = 128
序列长度
1
seq_len = 4096
层数
1
n_layers = 32

那么，单条请求的 KV Cache 显存占用为：


1
2
3
4
KV_Cache_Size = 2 × n_layers × seq_len × d_model × sizeof(dtype)

              = 2 × 32 × 4096 × 4096 × 2  (FP16)

              = 2,147,483,648 bytes

              ≈ 2 GB

当批量大小

1	batch_size = 8

时，仅 KV Cache 就需要约 16GB 显存。这还不包括模型权重、激活值和其他运行时开销。在 80GB 的 A100 上，这意味着单个推理实例几乎无法支撑更多并发。

更关键的是，KV Cache 的大小与序列长度呈线性关系。当大模型开始支持 32K、128K 甚至 1M 的上下文窗口时，KV Cache 的显存压力呈爆炸式增长，成为不可逾越的硬件天花板。

MHA 的 PyTorch 实现示例


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
import torch

import torch.nn as nn

import math



class MultiHeadAttention(nn.Module):

    def __init__(self, d_model, n_heads):

        super().__init__()

        self.n_heads = n_heads

        self.d_head = d_model // n_heads

        

        # MHA: 每个头都有独立的 Q, K, V 投影

        self.W_q = nn.Linear(d_model, d_model, bias=False)

        self.W_k = nn.Linear(d_model, d_model, bias=False)  # n_heads 个头的 K

        self.W_v = nn.Linear(d_model, d_model, bias=False)  # n_heads 个头的 V

        self.W_o = nn.Linear(d_model, d_model, bias=False)

    

    def forward(self, x, kv_cache=None):

        B, S, D = x.shape

        

        # 投影并分头

        Q = self.W_q(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)

        K = self.W_k(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)

        V = self.W_v(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)

        

        # KV Cache 更新：存储所有头的 K 和 V

        if kv_cache is not None:

            K_cached, V_cached = kv_cache

            K = torch.cat([K_cached, K], dim=2)  # 沿序列维度拼接

            V = torch.cat([V_cached, V], dim=2)

        

        # 注意力计算

        attn = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)

        attn = torch.softmax(attn, dim=-1)

        out = torch.matmul(attn, V)

        

        out = out.transpose(1, 2).contiguous().view(B, S, D)

        return self.W_o(out), (K, V)

注意上面的代码中，

和

的形状都是

1	(B, n_heads, S, d_head)

，每个头都存储了完整的 Key 和 Value 序列。这就是 KV Cache 显存膨胀的根源——每个头的 KV 都需要独立缓存。

MQA：最激进的 KV Cache 缩减方案

多查询注意力（Multi-Query Attention, MQA）由 Noam Shazeer 在 2019 年的论文《Fast Transformer Decoding: One Write-Head is All You Need》中首次提出。其核心思想极其简单粗暴：让所有 Query 头共享同一组 Key 和 Value 头。

在 MQA 中，无论模型有多少个 Query 头，Key 和 Value 都只有一个头。这意味着 KV Cache 的显存占用直接缩减为 MHA 的

1/n_heads

。对于 32 头的模型，这相当于 97% 的 KV Cache 显存节省。

MQA 架构详解

MQA 的修改集中在投影矩阵上：

1
W_q

维度不变：

1
(d_model, d_model)

，产生 n_heads 个 Query 头
1
W_k

维度缩减：

1
(d_model, d_head)

，只产生 1 个 Key 头
1
W_v

维度缩减：

1
(d_model, d_head)

，只产生 1 个 Value 头


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class MultiQueryAttention(nn.Module):

    def __init__(self, d_model, n_heads):

        super().__init__()

        self.n_heads = n_heads

        self.d_head = d_model // n_heads

        

        # MQA: Q 有 n_heads 个头, K 和 V 各只有 1 个头

        self.W_q = nn.Linear(d_model, d_model, bias=False)

        self.W_k = nn.Linear(d_model, self.d_head, bias=False)  # 单头 K

        self.W_v = nn.Linear(d_model, self.d_head, bias=False)  # 单头 V

        self.W_o = nn.Linear(d_model, d_model, bias=False)

    

    def forward(self, x, kv_cache=None):

        B, S, D = x.shape

        

        # Q: n_heads 个头

        Q = self.W_q(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)

        # K, V: 只有 1 个头，需要广播给所有 Q 头

        K = self.W_k(x).view(B, S, 1, self.d_head).transpose(1, 2)  # (B, 1, S, d_head)

        V = self.W_v(x).view(B, S, 1, self.d_head).transpose(1, 2)  # (B, 1, S, d_head)

        

        if kv_cache is not None:

            K_cached, V_cached = kv_cache

            K = torch.cat([K_cached, K], dim=2)

            V = torch.cat([V_cached, V], dim=2)

        

        # K 和 V 自动广播到所有 Q 头

        K_expanded = K.expand(-1, self.n_heads, -1, -1)

        V_expanded = V.expand(-1, self.n_heads, -1, -1)

        

        attn = torch.matmul(Q, K_expanded.transpose(-2, -1)) / math.sqrt(self.d_head)

        attn = torch.softmax(attn, dim=-1)

        out = torch.matmul(attn, V_expanded)

        

        out = out.transpose(1, 2).contiguous().view(B, S, D)

        return self.W_o(out), (K, V)  # KV Cache 只需存单头

MQA 的优势在于：

显存节省巨大：KV Cache 缩减为原来的 1/n_heads，32 头模型节省 97%
解码速度提升：每次生成新 token 时，只需从内存中读取一组 K 和 V，大幅减少内存带宽压力
实现简单：仅需修改投影矩阵维度，不改变注意力计算的核心逻辑

但 MQA 也有明显的代价：模型质量下降。Key 和 Value 只有一个头意味着所有 Query 头必须”共享视角”，注意力模式的表达能力大幅受限。在实际评测中，MQA 模型在困惑度（perplexity）上通常比同规模的 MHA 模型高 2-5%，在复杂推理任务上的差距更为明显。

模型架构对比示意图

GQA：在 MQA 和 MHA 之间寻找最优解

分组查询注意力（Grouped-Query Attention, GQA）由 Google Research 在 2023 年的论文《GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints》中提出。GQA 的核心思想是：将 Query 头分成若干组，每组共享一组 Key 和 Value 头。

可以认为 GQA 是 MHA 和 MQA 的统一泛化：

当
1
n_kv_heads = n_heads

时，GQA 退化为 MHA
当
1
n_kv_heads = 1

时，GQA 退化为 MQA
当
1
1 < n_kv_heads < n_heads

时，GQA 在两者之间取得平衡

这个看似简单的改动，在工程实践中展现出了惊人的效果。LLaMA-2-70B 使用 8 个 KV 头（

1	n_heads=64, n_kv_heads=8

），在保持与 MHA 模型相当的质量的同时，将 KV Cache 缩减了 87.5%。

GQA 的实现细节


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class GroupedQueryAttention(nn.Module):

    def __init__(self, d_model, n_heads, n_kv_heads):

        super().__init__()

        self.n_heads = n_heads

        self.n_kv_heads = n_kv_heads

        self.n_groups = n_heads // n_kv_heads  # 每组包含的 Q 头数

        self.d_head = d_model // n_heads

        

        # GQA: Q 有 n_heads 个头, K 和 V 有 n_kv_heads 个头

        self.W_q = nn.Linear(d_model, n_heads * self.d_head, bias=False)

        self.W_k = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)

        self.W_v = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)

        self.W_o = nn.Linear(n_heads * self.d_head, d_model, bias=False)

    

    def forward(self, x, kv_cache=None):

        B, S, D = x.shape

        

        Q = self.W_q(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)

        K = self.W_k(x).view(B, S, self.n_kv_heads, self.d_head).transpose(1, 2)

        V = self.W_v(x).view(B, S, self.n_kv_heads, self.d_head).transpose(1, 2)

        

        if kv_cache is not None:

            K_cached, V_cached = kv_cache

            K = torch.cat([K_cached, K], dim=2)

            V = torch.cat([V_cached, V], dim=2)

        

        # 关键操作：将 K 和 V 从 n_kv_heads 扩展到 n_heads

        # repeat_interleave — 每组内连续重复

        K = K.repeat_interleave(self.n_groups, dim=1)

        V = V.repeat_interleave(self.n_groups, dim=1)

        

        attn = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)

        attn = torch.softmax(attn, dim=-1)

        out = torch.matmul(attn, V)

        

        out = out.transpose(1, 2).contiguous().view(B, S, -1)

        return self.W_o(out), (K, V)

这里最关键的操作是

1	repeat_interleave

。它将

1	n_kv_heads

个 KV 头扩展为

n_heads

个，每个 KV 头被连续复制

n_groups

次。这意味着第 0 组的 Q 头（头 0 到头

1	n_groups-1

）都关注同一个 KV 头（头 0），第 1 组的 Q 头关注同一个 KV 头（头 1），以此类推。

GQA 的 KV Cache 显存计算

GQA 模型中，KV Cache 的显存占用公式为：


1
GQA_KV_Cache = 2 × n_layers × seq_len × n_kv_heads × d_head × sizeof(dtype)

与 MHA 的比例为：


1
ratio = n_kv_heads / n_heads

模型	n_heads	n_kv_heads	KV Cache 比例	节省比例
MHA (LLaMA-1-65B)	64	64	100%	0%
GQA (LLaMA-2-70B)	64	8	12.5%	87.5%
GQA (Mistral-7B)	32	8	25%	75%
GQA (Llama-3-8B)	32	8	25%	75%
MQA (Falcon-40B)	128	1	0.78%	99.2%

从 MHA 预训练模型转换到 GQA 的工程实践

GQA 论文最有工程价值的贡献之一是提出了一种从已有 MHA 模型转换到 GQA 模型的方法，而不需要从头训练。这意味着我们可以复用已有的预训练权重，大幅降低训练成本。

转换的核心思路是”均值池化”（mean averaging）：

将 MHA 模型中属于同一组的 Q 头对应的 K 和 V 权重取平均
用这个均值作为 GQA 模型中该组的 KV 头权重
进行少量预训练式微调（pre-training fine-tuning）恢复模型质量


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def convert_mha_to_gqa(state_dict, n_heads, n_kv_heads):

    """将 MHA 权重转换为 GQA 权重

    

    Args:

        state_dict: MHA 模型的状态字典

        n_heads: 原始 Q 头数

        n_kv_heads: 目标 KV 头数

    """

    n_groups = n_heads // n_kv_heads

    new_state_dict = {}

    

    for key, value in state_dict.items():

        if '.W_k.' in key or '.W_v.' in key:

            d_model = value.shape[0]

            d_head = value.shape[1] // n_heads

            reshaped = value.view(d_model, n_heads, d_head)

            

            # 分组并取均值

            grouped = reshaped.view(d_model, n_kv_heads, n_groups, d_head)

            averaged = grouped.mean(dim=2)

            

            new_state_dict[key] = averaged.reshape(d_model, n_kv_heads * d_head)

        else:

            new_state_dict[key] = value

    

    return new_state_dict

论文实验表明，从 LLaMA-2 的 MHA 检查点转换到 GQA-8（8 个 KV 头），仅需约 原始预训练 5% 的训练量即可恢复模型质量。这是一个极其实用的结果——它让已经在 MHA 架构上投入数百万美元训练成本的团队，能够以极低的代价迁移到 GQA 架构。

vLLM 中 GQA 的实现优化：融合与排布

在实际推理引擎（如 vLLM、TensorRT-LLM）中，GQA 的实现远比上述学术示例复杂。核心优化集中在两个方面：内存排布优化和计算融合。

内存排布优化

在 vLLM 的 PagedAttention 实现中，KV Cache 被组织成固定大小的”页”（page），每个页存储若干 token 的 KV 向量。对于 GQA 模型，每个页中的 KV 向量数量为

1	n_kv_heads × d_head

，而非 MHA 的

1	n_heads × d_head

。这直接意味着：

每个 KV 页的物理大小缩减为 MHA 的
1
n_kv_heads / n_heads
相同显存下可以分配更多页，支持更多并发请求
页表（block table）的条目数不变，但每条目对应的物理块更紧凑

计算融合：避免显式的 repeat_interleave

在上述 PyTorch 示例中，我们使用

1	repeat_interleave

将 KV 头扩展到与 Q 头相同的数量。但在 GPU 内核中，这个操作会引入不必要的显存读写。vLLM 和 TensorRT-LLM 的做法是将 KV 头扩展逻辑融合到注意力 kernel 内部：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
// 伪代码: CUDA kernel 中的 GQA 注意力计算

// 无需显式扩展 K/V, 而是在计算时按组索引



__global__ void gqa_attention_kernel(

    float* Q, float* K, float* V, float* O,

    int n_heads, int n_kv_heads, int d_head, int seq_len

) {

    int q_head_idx = blockIdx.x;  // 当前 Q 头索引

    int kv_head_idx = q_head_idx / (n_heads / n_kv_heads);  // 映射到 KV 头

    

    // 直接从对应的 KV 头读取, 避免显式复制

    for (int s = 0; s &lt; seq_len; s++) {

        float score = 0.0f;

        for (int d = 0; d &lt; d_head; d++) {

            score += Q[q_head_idx * seq_len * d_head + d]

                   * K[kv_head_idx * seq_len * d_head + d];

        }

        score /= sqrtf((float)d_head);

        // softmax + weighted sum ...

    }

}

这种融合策略完全消除了 KV 头扩展的显存开销，同时减少了全局内存访问次数。在 A100 上，融合实现比先扩展再计算的方案快 15-25%。

不同 KV Head 配置下的性能与质量权衡

选择多少个 KV Head 并非简单的”越多越好”或”越少越好”。GQA 论文通过系统的实验揭示了

1	n_kv_heads

对推理速度和模型质量的影响模式。

推理吞吐量对比

在相同的硬件条件下，不同的 KV Head 配置对推理吞吐量的影响可以通过以下公式估算：


1
2
3
4
5
6
7
# 假设模型参数量固定, 显存带宽是瓶颈

# 解码阶段的延迟主要由 KV Cache 读取决定



decoding_latency ∝ n_kv_heads × seq_len × d_head



# 吞吐量 (tokens/sec) 反比于延迟

throughput ∝ n_heads / n_kv_heads

实测数据（A100-80GB, LLaMA-2-70B 系列, batch_size=8, seq_len=2048）：

配置	n_kv_heads	吞吐量 (tok/s)	PPL (↓)
MHA	64	28	3.32
GQA-16	16	42	3.35
GQA-8	8	58	3.38
GQA-4	4	72	3.47
MQA	1	95	3.61

从表中可以看出，GQA-8 在吞吐量上比 MHA 提升了约 107%，而 PPL 仅增加 0.06（约 1.8%），是性价比最高的配置点。这正是 LLaMA-2-70B 和 Llama-3 系列选择 8 个 KV Head 的核心原因。

主流开源模型的 KV Head 配置一览

了解各模型的 KV Head 配置，有助于在实际部署中预估显存需求和推理性能：

模型	参数量	n_heads	n_kv_heads	架构	上下文长度
LLaMA-1-65B	65B	64	64	MHA	2048
LLaMA-2-70B	70B	64	8	GQA	4096
Llama-3-8B	8B	32	8	GQA	8192
Llama-3-70B	70B	64	8	GQA	8192
Llama-3.1-405B	405B	128	8	GQA	128K
Mistral-7B-v0.1	7B	32	8	GQA	8192
Mixtral-8x7B	47B	32	8	GQA	32K
Falcon-180B	180B	232	8	GQA	2048
Qwen2-72B	72B	64	8	GQA	128K
DeepSeek-V2	236B	128	128	MLA*	128K

*DeepSeek-V2 使用的是 MLA（Multi-head Latent Attention），本质上是对 KV Cache 的低秩压缩，与 GQA 的思路不同但目标一致。

一个明显的趋势是：8 个 KV Head 已经成为大模型的事实标准。无论是 Meta 的 Llama 系列还是 Mistral 系列，70B 级别模型几乎都选择了 8 个 KV Head。这个数字在质量与效率之间找到了最佳平衡点。

实战：如何选择合适的 KV Head 数量

如果你正在训练自己的模型或对开源模型进行定制化推理部署，以下决策流程可以帮助你选择合适的 KV Head 数量：

场景 1：从头训练新模型

对于 7B-13B 规模的模型，推荐配置：

通用场景：
1
n_kv_heads = n_heads / 4

，例如 32 头模型用 8 个 KV 头
推理优先：
1
n_kv_heads = n_heads / 8

，牺牲少量质量换取更快推理
质量优先：
1
n_kv_heads = n_heads / 2

，适合需要复杂推理的任务

场景 2：已有 MHA 模型的推理优化

如果你已经有一个 MHA 模型，不想重新训练，可以通过以下方式优化 KV Cache：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 方法1: 使用 GQA 转换 + 少量微调 (推荐)

# 1. 转换权重

# 2. 用原始预训练数据的 5% 进行继续训练

# 3. 约 2000-5000 步即可恢复质量



# 方法2: KV Cache 量化 (无需重训练)

# 将 FP16 的 KV Cache 量化为 FP8 或 INT8

# 精度损失极小 (~0.1% PPL), 显存节省 50%



import torch



def quantize_kv_cache_fp8(K, V):

    """将 KV Cache 从 FP16 量化为 FP8 (E4M3)"""

    K_fp8 = K.to(torch.float8_e4m3fn)

    V_fp8 = V.to(torch.float8_e4m3fn)

    return K_fp8, V_fp8



# 方法3: KV Cache 分页 + 智能卸载

# 将不活跃的 KV 页卸载到 CPU 内存

# 仅在需要时加载回 GPU, 适合长上下文场景

场景 3：长上下文推理优化

当上下文长度超过 32K 时，即使使用 GQA，KV Cache 的显存压力仍然很大。此时需要结合多种优化策略：

GQA + KV Cache 量化：双重压缩，8 KV Head + FP8 量化可将 KV Cache 压缩至 MHA 的 6.25%
Sliding Window Attention：Mistral 使用的策略，只缓存最近 W 个 token 的 KV，配合滚动窗口
Token 蒸馏/剪枝：在预填充阶段识别不重要的 token，不缓存其 KV 向量


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# Sliding Window + GQA 组合示例

class SlidingWindowGQA(nn.Module):

    def __init__(self, d_model, n_heads, n_kv_heads, window_size=4096):

        super().__init__()

        self.gqa = GroupedQueryAttention(d_model, n_heads, n_kv_heads)

        self.window_size = window_size

    

    def forward(self, x, kv_cache=None):

        out, (K, V) = self.gqa(x, kv_cache)

        

        # 滑动窗口裁剪: 只保留最近 window_size 个 token

        if K.shape[2] &gt; self.window_size:

            K = K[:, :, -self.window_size:, :]

            V = V[:, :, -self.window_size:, :]

        

        return out, (K, V)

前沿方向：MLA 与动态 KV 压缩

GQA 和 MQA 通过减少 KV Head 数量来压缩 KV Cache，本质上是一种”结构化压缩”——压缩发生在模型架构设计阶段。但还有另一条路径：对 KV 向量本身进行低秩压缩，代表方案是 DeepSeek-V2 的 MLA（Multi-head Latent Attention）。

MLA 的核心思路是：不再直接缓存 Key 和 Value 向量，而是缓存一个低维的”潜在向量”（latent vector），在需要时通过上投影矩阵恢复出完整的 KV 向量。这样做的好处是：

压缩比更高：可以将 KV Cache 压缩到 MHA 的 5% 以下
不损失注意力头的多样性：每个 Q 头仍然看到不同的 KV 表示
适合超长上下文：128K 甚至更长序列下的 KV Cache 管理更加可行

不过 MLA 的实现复杂度显著高于 GQA，且对 RoPE 位置编码的兼容性需要额外处理（需要将 RoPE 解耦为独立的位置感知分量）。目前 MLA 主要在 DeepSeek 系列模型中使用，尚未成为业界通用方案。

另一个值得关注的方向是动态 KV 压缩——根据输入序列的特性，在运行时决定哪些 token 的 KV 需要保留、哪些可以丢弃。这与传统注意力稀疏化（如 Longformer、BigBird）的思路一脉相承，但更侧重于推理阶段的 KV Cache 管理而非训练阶段的注意力模式设计。典型的代表包括 H2O（Heavy-Hitter Oracle）和 Scissorhands 等工作。

总结与最佳实践

从 MHA 到 GQA 到 MQA 的演进，本质上是在模型表达能力和推理效率之间寻找最优折中。以下是我们基于大量实践总结的推荐方案：

新训练模型：直接采用 GQA 架构，
1
n_kv_heads

选择 8 是目前最稳妥的选择
已有 MHA 模型：通过 GQA 转换 + 5% 预训练量微调，以最小成本获得推理加速
极致推理性能：MQA 配合 KV Cache 量化，适合对延迟极其敏感的在线服务
超长上下文：GQA + Sliding Window + KV 量化的组合拳，必要时考虑 MLA

KV Cache 优化是大模型推理工程的核心课题。理解 MHA/GQA/MQA 的设计原理和工程权衡，不仅能帮助你做出正确的架构选择，更能让你在面对推理性能瓶颈时有的放矢。在模型规模持续增长、上下文窗口不断扩展的今天，KV Cache 优化的重要性只会越来越高。