Attention

Attention 让当前 token 根据 Query 与历史 Key 的相关性，对历史 Value 做加权聚合。本文只介绍所有模型共享的基础计算，以及 MHA、MQA、GQA 三种 KV Head 共享方式；MLA、DSA、CSA、HCA、GDN 和 KDA 等模型专属机制分别放在对应模型文章中。

Scaled Dot-Product Attention

设输入为 $X\in {\mathbb{R}}^{B\times N\times D}$，其中 $B$ 是 Batch Size，$N$ 是 Sequence Length，$D$ 是 Hidden Size。线性投影生成 Query、Key 和 Value：

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V$$

单个 Attention Head 的输出为：

$$\mathrm{Attention}(Q,K,V)=\mathrm{Softmax}(\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_h}}+M)V$$

$d_h$ 是 Head Dimension，$M$ 是 Attention Mask。Decoder 使用 Causal Mask，使位置 $t$ 只能读取自身及其之前的位置。

Hidden States
  ├─ Wq → Query ───────────────┐
  ├─ Wk → Key   → KV Cache ────┼─→ QKᵀ / √dₕ → Mask → Softmax
  └─ Wv → Value → KV Cache ────┘                    │
                                                     └─→ × Value → Output

Multi-Head Attention

Multi-Head Attention 将隐藏维度切成 $h_q$ 个 Query Heads。不同 Head 可以学习不同的匹配模式，所有 Head 的结果拼接后再经过输出投影：

$$\mathrm{MHA}(X)=W_O\mathrm{Concat}({\mathrm{head}}_1,\dots ,{\mathrm{head}}_{h_q})$$$${\mathrm{head}}_i=\mathrm{Attention}(Q_i,K_{g(i)},V_{g(i)})$$

$g(i)$ 表示第 $i$ 个 Query Head 使用哪一个 KV Head。MHA、MQA 和 GQA 的主要区别就是映射 $g(i)$。

KV Head Sharing

令 Query Head 数量为 $h_q$，KV Head 数量为 $h_{kv}$。三种结构执行相同的 Softmax Attention，只改变多少个 Query Heads 共享一组 Key 和 Value。

MHA

Multi-Head Attention（MHA）为每个 Query Head 配置独立的 KV Head：

$$h_{kv}=h_q,g(i)=i$$

MHA 的 KV 表达能力最完整，但 KV Cache 宽度也最大。

MQA

Multi-Query Attention（MQA）让全部 Query Heads 共享唯一的 KV Head：

$$h_{kv}=1,g(i)=0$$

MQA 显著减少 KV Cache 和 Decode 阶段读取的 K/V 数据量，但所有 Query Heads 只能从同一组 K/V 表示中读取信息。

GQA

Grouped-Query Attention（GQA）位于 MHA 与 MQA 之间。每组 $r=h_q/h_{kv}$ 个 Query Heads 共享一个 KV Head：

$$1<h_{kv}<h_q,g(i)=\lfloor \frac{i}{r}\rfloor$$

GQA 用少量 KV Heads 保留比 MQA 更丰富的 K/V 表示，同时获得大部分 KV Cache 与内存带宽收益。

类型	KV Heads	共享关系	主要权衡
MHA	$h_{kv}=h_q$	每个 Query Head 独占 K/V	表达能力强，KV Cache 最大
GQA	$1<h_{kv}<h_q$	一组 Query Heads 共享 K/V	表达能力与推理成本折中
MQA	$h_{kv}=1$	所有 Query Heads 共享 K/V	KV Cache 最小，共享程度最高

📄 Transformers：GQA/MQA 的逻辑 Head 展开

def repeat_kv(hidden_states, n_rep):
    batch, num_key_value_heads, seq_len, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(
        batch, num_key_value_heads, n_rep, seq_len, head_dim
    )
    return hidden_states.reshape(
        batch, num_key_value_heads * n_rep, seq_len, head_dim
    )


def eager_attention_forward(module, query, key, value, attention_mask, scaling, ...):
    key = repeat_kv(key, module.num_key_value_groups)
    value = repeat_kv(value, module.num_key_value_groups)
    scores = torch.matmul(query, key.transpose(2, 3)) * scaling
    scores = scores + attention_mask
    weights = torch.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype)
    return torch.matmul(weights, value), weights

repeat_kv 表达逻辑共享关系。支持 GQA/MQA 的高效 Kernel 可以直接处理不同数量的 Query Heads 与 KV Heads，不一定在显存中真正复制 K/V。

推理成本

每个缓存元素占 $s$ Byte 时，标准逐 token KV Cache 的单层容量为：

$$\mathrm{KVCache}=2BN{h}_{kv}{d}_{h}s$$

系数 $2$ 分别对应 Key 和 Value。减少 $h_{kv}$ 会线性降低 KV Cache 容量和 Decode 的 K/V 读取量，但不会减少 Query Head 数量，也不会改变 Query 访问多少个历史 token。

对于读取全部历史的 Full Attention：

$${\mathrm{Cost}}_{\text{prefill}}=O(N^2{h}_q{d}_h)$$$${\mathrm{Cost}}_{\text{decode/token}}=O(N{h}_q{d}_h)$$

因此，MHA → GQA → MQA 主要优化 KV Cache 与内存带宽；要减少随上下文长度增长的 Attention 计算，还需要改变 KV 表示或访问模式。DeepSeek V3 的 MLA、DeepSeek V3.2 的 DSA、DeepSeek V4 的 CSA/HCA 和 Kimi K3 的 KDA 分别采用了不同路线。

FlashAttention 不改变 MHA、GQA 或 MQA 的语义。它通过分块计算和减少 HBM 读写优化执行过程，仍然属于对相同 Attention 结果的 Kernel 实现。

参考

1. Attention Is All You Need
2. Fast Transformer Decoding: One Write-Head is All You Need
3. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
4. Transformers：Qwen2 Attention 实现