GQA / MQA / MLA 对 KV cache 占用与解码带宽的影响

2025-12-16

在同一组 **** 的 ****（8 卡，单卡 64GB HBM，卡间互联 392 GB/s）上并排起两个模型服务：一个注意力结构是 GQA，另一个是 MLA，两者参数量同量级、层数接近，用同样的 --max-model-len 启动。启动日志里 KV 池的可用块数相差一个数量级以上；压到长上下文时，GQA 那个服务的 decode 受显存带宽约束、ITL p50 上抬，MLA 那个服务的 KV 驻留明显更小。把一条现成的 KV 占用估算式直接套到两个模型上，套用后结果与日志不符，原因是式中的 KV 头数 H_kv 与单头维度 D 在两种结构下的口径不同。要分清的是：三种结构下这两个量各自怎么取，占用与解码带宽差在哪里。

三种结构下 KV 头数与 query 头数的关系

设 query 头数为 H_q。三种结构对 KV 的存法如下：

MHA：每个 query 头各自配一套 K、V，H_kv = H_q，KV cache 最大。
GQA：把 query 头分成若干组，每组共享一套 K、V，H_kv = H_q / group_size，H_kv 远小于 H_q。
MQA：全部 query 头共享一套 K、V，H_kv = 1，KV cache 最小，但表达力受限于单一 KV。
MLA：不按头存 K/V，而是把每条 token 的 K/V 压缩成一个潜向量（latent）再加一小段解耦的 RoPE 分量，所有头共用这一份缓存。

结构	KV 头数与 query 头数	每层每 token 缓存的量	缓存形状口径
MHA	`H_kv = H_q`	`2 × H_q × D`	每头独立的 K、V
GQA	`H_kv = H_q / group_size`	`2 × H_kv × D`	每头独立的 K、V，头数已归约
MQA	`H_kv = 1`	`2 × 1 × D`	单套 K、V
MLA	不适用（潜向量被所有头共享）	约 `(kv_lora_rank + qk_rope_head_dim) × b`	`[num_pages, page_size, kv_lora_rank + qk_rope_head_dim]`，decode 下 KV 头数恒为 1

占用与解码带宽的口径不能混用

沿用 KV cache 占用估算式：

整个模型 · 单序列 = 2 × L × H_kv × D × S × b
每卡（KV 头按 TP 切分） = 2 × L × (H_kv / TP) × D × S × b × N
# L 注意力层数，S 单序列 token 数，b 每元素字节数，N 并发，TP tensor_parallel_size

对 MHA / MQA / GQA，H_kv 直接取缓存里实际的 KV 头数，D 取单头维度，这式子可用。取一组常见规格作估算（非实测）：L=32、H_kv=8、D=128、b=2，单序列到 8192 token 时 2 × 32 × 8 × 128 × 8192 × 2 约为 1 GiB；把 H_kv 改成 1（MQA 形态），同一序列降到约 128 MiB，差出约 8 倍。若换成 H_kv = H_q = 64 的 MHA，则升到约 8 GiB。层数不变，仅头数的差别就跨了一个量级。

对 MLA，这式子不能直接用。MLA 缓存的是压缩后的潜向量，不按头分开存 K、V。它的 KV 缓存张量形状是 [num_pages, page_size, kv_lora_rank + qk_rope_head_dim]，其中 KV 头数在 decode 阶段恒为 1：潜向量被所有头共享，缓存里不再按 KV 头数分开，页粒度退化为 page_size=1 时按原始表示存，占用最浪费。

# MLA 的 KV 缓存张量形状（示意，取值随模型配置）
kv_buffer: [num_pages, page_size, kv_lora_rank + qk_rope_head_dim]
# 例：kv_lora_rank=512, qk_rope_head_dim=64 → 每 token 每层缓存宽度 576 元素
# 对照 GQA：每 token 每层 2 × H_kv × D（例 H_kv=8, D=128 → 2048 元素）

取一组常见取值作估算（非实测）：kv_lora_rank=512、qk_rope_head_dim=64，则每层每 token 的缓存宽度是 512 + 64 = 576 个元素，bf16（b=2）下为 1152 字节。估算时应把式中的 H_kv 与 D 换成按这个实际张量形状取的等效值，不能沿用 2 × H_kv × D。qk_nope_head_dim、qk_rope_head_dim、v_head_dim 与 kv_lora_rank 这些键的取值随模型配置给出，估算前先读模型自己的配置文件，不要套用别的 MLA 模型的数。

解码带宽同理。decode 每步要回读整条序列累积的 KV。GQA/MQA 回读的是归约后的 K、V，H_kv 越小，每步搬运的字节越少，但单头承载的信息越集中。MLA 回读的是压缩后的潜向量，并要按头展开参与注意力计算，它的显存访问形态与按头存的 K、V 不同。判定 decode 是否受 KV 带宽约束时，看的是 ITL 与 TTFT 的相对越界：ITL p50 超过 SLO 上限而 TTFT p50 未超，说明瓶颈落在 decode 稳态的 KV 回读；num_key_value_heads 很小（GQA/MQA 形态）或 is_mla() 时，这一判定更可能成立。

现象

同规模模型的显存需求差异巨大。 参数量与层数接近，KV 池可用块数却差一个量级以上，服务能起的 max-model-len 与并发上限随之不同。
长上下文表现不同。 长序列下 KV 占用按 S 线性增长，H_kv 或潜向量维度小的结构能容纳更长序列或更高并发；反之先受并发约束触顶。
同一估算式套不同结构算不准。 直接代入注意力头数当 H_kv，对 GQA 会高估、对 MLA 会错口径，得到的需求远大于或小于实际分配。
decode 侧的带宽瓶颈形态。 服务能起、单条可跑，并发上去后 ITL p95 抖动而 TTFT 正常，属于 KV 回读/带宽受限，不是 prefill 算力受限。

触发条件

结构	典型现象	先受约束的方面
MHA	同规模下 KV 池块数最少	显存容量最先触顶，长上下文或高并发先起不来
GQA	`H_kv` 随分组归约，占用按比例下降	长上下文更从容，单组 KV 承载多个 query 头
MQA	`H_kv=1`，占用最小	表达力受单一 KV 限制，质量敏感任务需实测
MLA	缓存压缩为潜向量，按页存放	decode 每步回读潜向量，对显存带宽敏感

MLA 结构还多一层约束：它的注意力走独立的后端族，is_mla() 与标准注意力的实现、对 head_size 与 KV dtype 的要求都不同。部分加速卡的首代型号没有生产级 MLA decode 内核，只能回退到通用实现，此时「用了 MLA 结构」不等于「拿到了 MLA 的带宽收益」，判定前先确认实际命中的后端族。

定位步骤

从模型配置读结构。 读 num_attention_heads（H_q）与 num_key_value_heads（H_kv）；H_kv = H_q 为 MHA，H_kv = 1 为 MQA，介于二者为 GQA；存在 qk_nope_head_dim / kv_lora_rank 一类键则为 MLA。
按结构取口径算占用。 非 MLA 结构代入 H_kv 与 D；MLA 结构按潜向量张量形状取等效值，不代入注意力头数。
与启动日志的 KV 池分配对照。 读日志里分配的块数与每块字节数，与手算值对比，差额即激活与工作区占用；对不上时先核对第 1、2 步的口径。
分段测 decode 带宽。 固定 prompt 长度、逐步抬并发，观察 ITL 与 HBM 带宽利用率曲线；ITL 随并发恶化而带宽贴顶，即为带宽受限。
确认 attention 后端命中。 MLA 与标准注意力走不同的后端族（is_mla()），head_size、num_kv_heads 等特征决定目标后端的 guard 是否满足，失格时框架回退到另一族实现。

记录字段

# KV 结构与占用结论记录清单（字段名示意）
model_name:               # 模型
attention_family:         # mha / gqa / mqa / mla（is_mla()）
num_attention_heads:      # H_q
num_key_value_heads:      # H_kv
head_dim:                 # D
mha_mla_keys:             # MLA 专用键：qk_nope_head_dim / qk_rope_head_dim / v_head_dim / kv_lora_rank
num_layers:               # L
kv_dtype:                 # bf16 / fp8，决定 b
tp_size:                  # TP
context_len:              # S 上限，max_model_len
concurrency:              # N
kv_pool_blocks:           # 启动日志中的 KV 池可用块数
itl_p50_ms / ttft_p50_ms: # 判定 decode 侧带宽瓶颈
attention_backend:        # 实际命中的后端
verdict:                  # 带宽受限 / 显存受限 / 口径待核对

小结

H_kv 表示每层缓存几套 K、V：MHA 下等于 query 头数，GQA 下是归约后的组数，MQA 下为 1，MLA 下这套计数不适用、改按潜向量形状取值。占用估算式只在非 MLA 结构下可直接代入 H_kv 与 D；层数不变时，仅 H_kv 的差别就能让占用与解码带宽跨一个量级。套用估算式前先读模型配置确定结构，MLA 尤其要按实际缓存的张量形状取等效维度，判定瓶颈时用 ITL 与 TTFT 的相对越界区分 decode 带宽与 prefill 算力。