【AI 工业核题 D3】GQA / MQA 分组查询注意力（共享 KV 头）（Grouped-Query & Multi-Query Attention）深度实现与原理解析

题目分类：Part D · 注意力机制与 Transformer 核心组件 (Part D · Attention Mechanisms & Transformer Blocks) | 难度等级：Medium | 工业重要度：工业基石 (核心高频)

一、核心题意与背景

大幅降低自回归解码阶段 KV Cache 显存占用的关键创新，MHA 到 MQA 的完美折中。

Industrial-grade implementation and mathematical foundations of Grouped-Query & Multi-Query Attention.

二、数学原理与公式推导

KV Cache 带宽墙与架构演进

在大模型推理生成时，瓶颈不在于算力（FLOPS），而在于显存访存带宽（Memory Bandwidth Bound）。每次自回归生成 1 个 Token，都要将全量上下文的 KV Cache 从 HBM 读到 SRAM。
– MHA（标准多头）：$H_Q = H_{KV}$，每个 Q 头拥有独立的 KV 头，显存占用大；
– MQA（Multi-Query Attention）：$H_{KV} = 1$，所有 Q 头共享同一个单一的 Key/Value 头，显存降到最低，但表达能力略有折损；
– GQA（Grouped-Query Attention）：将 $H_Q$ 个 Query 头均匀划分为 $G$ 个组，每个组内共享一个 KV 头（即 $H_{KV} = G$）。
在 LLaMA-2-70B、LLaMA-3（$H_Q=32, H_{KV}=8$）中，KV Cache 显存直接缩减为原先的 $frac{1}{4}$，同时性能无损。

三、工业级 Python 核心实现

import numpy as np

def repeat_kv(x: np.ndarray, num_rep: int) -> np.ndarray:
    """
    将 KV 头重复展开以对齐 Query 头的数量。
    x 形状: (B, H_kv, S, D_k)
    返回: (B, H_q, S, D_k)，其中 H_q = H_kv * num_rep
    """
    if num_rep == 1:
        return x
    B, H_kv, S, D_k = x.shape
    # 增加扩展维度并广播
    x = x[:, :, np.newaxis, :, :]  # (B, H_kv, 1, S, D_k)
    x = np.repeat(x, num_rep, axis=2)  # (B, H_kv, num_rep, S, D_k)
    return x.reshape(B, H_kv * num_rep, S, D_k)

def grouped_query_attention(
    q: np.ndarray,      # (B, H_q, S, D_k)
    k: np.ndarray,      # (B, H_kv, S, D_k)
    v: np.ndarray,      # (B, H_kv, S, D_k)
    is_causal: bool = True
) -> np.ndarray:
    B, H_q, S, D_k = q.shape
    H_kv = k.shape[1]
    assert H_q % H_kv == 0, "Query 头数必须是 KV 头数的整数倍"
    num_rep = H_q // H_kv

    # 将 KV 广播扩展为与 Q 头数相同
    k_expanded = repeat_kv(k, num_rep)  # (B, H_q, S, D_k)
    v_expanded = repeat_kv(v, num_rep)  # (B, H_q, S, D_k)

    # 执行常规多头注意力点积
    scores = np.matmul(q, k_expanded.swapaxes(-1, -2)) / np.sqrt(D_k)
    if is_causal:
        mask = np.triu(np.ones((S, S), dtype=bool), k=1)
        scores = np.where(mask, -1e9, scores)

    scores_max = np.max(scores, axis=-1, keepdims=True)
    probs = np.exp(scores - scores_max)
    probs /= np.sum(probs, axis=-1, keepdims=True)
    return np.matmul(probs, v_expanded)

四、自动化单元测试与边界断言

import numpy as np
B, S, D_k = 2, 4, 8
H_q = 8
H_kv = 2  # 4 个 Q 头共享 1 个 KV 头
q = np.random.randn(B, H_q, S, D_k)
k = np.random.randn(B, H_kv, S, D_k)
v = np.random.randn(B, H_kv, S, D_k)
out = grouped_query_attention(q, k, v)
assert out.shape == (B, H_q, S, D_k), "输出形状不正确"
print("✓ GQA 分组查询注意力自测通过")

五、张量形状与维度变换流 (Tensor Flow)

• 中文解析：k: (B, H_kv, S, D_k) -> repeat_kv -> (B, H_q, S, D_k) -> 与 q: (B, H_q, S, D_k) 点积 -> (B, H_q, S, D_k)
• 英文对齐：k: (B, H_kv, S, D_k) -> repeat_kv -> (B, H_q, S, D_k) -> 与 q: (B, H_q, S, D_k) 点积 -> (B, H_q, S, D_k)

六、工业级数值稳定性避坑清单 (Checklist)

• ⚠️ 在实际底层 CUDA / Triton 实现中，无需物理上执行 repeat 显存拷贝，直接在读取指针寻址时做模运算索引即可
• ⚠️ KV Cache 存储时只需存储原始的 H_kv 形状，显存占用严格减少为 H_kv / H_q

English Checklist:
– Ensure proper multi-dimensional tensor broadcasting and keepdims retention.
– Enforce numerical guards (eps clamping and overflow thresholds) during exponentiation and division.
– Verify train versus eval mode behavioral distinctions (e.g. frozen running statistics and dropout bypass).

七、考场秒记心法口诀

💡 头多显存吃不消，多组共享一组刀，广播对齐算点积，显存直缩四分之一

Master Grouped-Query & Multi-Query Attention: enforce numerical stability, check tensor shapes, and eliminate redundant memory allocations.

八、高频面试追问与答题策略

Q1：DeepSeek-V2 / V3 提出的 MLA（Multi-Head Latent Attention）相比 GQA 又带来了什么突破？
(EN: What are the key trade-offs and memory bottlenecks when deploying Grouped-Query & Multi-Query Attention in high-throughput inference?)

答：GQA 是直接在头维度上做简单共享，仍需存储完整精度的特征；MLA 则是使用低秩矩阵分解将 Key 和 Value 联合压缩为一个极小的隐层潜变量（Latent Vector），并在解压前结合解耦 RoPE，将 KV Cache 压缩比推至前所未有的极限（约为原生 MHA 的 1/10 显存）。

(EN: Memory bandwidth (HBM to SRAM I/O) is the primary latency factor. Fusing element-wise operations and avoiding intermediate tensor materialization significantly outperforms naive implementations.)