先固定一组参数。我们用它把单头 Attention、多头、KV Cache 和分组共享完整算一遍。
| 符号 | 含义 | 本文取值 |
|---|---|---|
| \(D\) | 输入 hidden size | 4096 |
| \(H_q\) | Q head 数 | 32 |
| \(H_{kv}\) | KV head 数 | MHA 为 32，GQA 为 8，MQA 为 1 |
| \(d_h\) | 每个 head 的 Q/K/V 维度 | 128 |
| \(S_q\) | 当前计算的 query 数 | Prefill 示例为 4096，Decode 为 1 |
| \(T\) | 当前 query 可见的历史长度，包含当前位置 | Decode 示例为 8192 |
| \(b\) | 一个缓存元素的字节数 | BF16 为 2 |
为便于读公式，先省略 batch 轴，所有 token 都写成行向量。Q/K 可以先做位置编码，下面的 Attention 公式使用位置处理后的 Q/K。
1. 一个 head 到底计算什么
给一个输入 \(X\in\mathbb R^{S_q\times D}\)。单个 head 使用三个投影矩阵，把输入变成 query、key、value。首次 Prefill 中 Q/K/V 都由同一段输入生成：
所以 Q、K、V 的形状都是 \(S_q\times d_h\)。\(W_Q\) 学习怎样提问，\(W_K\) 学习哪些特征便于匹配，\(W_V\) 学习匹配到某个位置后提取哪些信息。三份权重分别学习，结果一般不同。
取第 \(i\) 个 query。它和第 \(j\) 个 key 的匹配分数为
这里的求和消去了 head 维度，留下 query 位置和 key 位置。因果 mask \(M_{ij}\) 在允许访问的位置为 0，在未来位置为负无穷。对每个 query 的整行分数做 Softmax：
然后对 value 加权求和：
第一次求和沿特征维 \(d_h\)，计算“匹配程度”；第二次沿位置维 \(T\)，汇总历史信息。写成矩阵乘就是
两个 GEMM 的规约维不同。QK 的 K 维是 head dimension，PV 的 K 维是历史长度。后面看 NPU 阵列时要一直保留这个区别。标准缩放点积 Attention 和多头组合见 Transformer 原论文。Attention Is All You Need
2. 用两个位置手算一次
假设缩放、mask 后的分数为 \([2,0]\)。为了简化，只看 V 的一个通道，两个位置的值为 \([10,20]\)。那么
换一个 query，分数变成 \([0,2]\)，同一份 K/V 就得到
这说明同一份历史 K/V 可以被不同 query 以不同方式读取。MQA/GQA 就会利用这一点。
3. MHA 怎样把多个 head 放在一起
32 个 head 各有一套投影。实现中通常把它们的权重横向拼接，一次 GEMM 生成所有 head：
得到 \(Q_{\rm flat}\in\mathbb R^{S_q\times4096}\) 后，把最后一维拆成 \([32,128]\)，再按内核需要排列为 \([32,S_q,128]\)。K/V 同理。每个 head 独立执行刚才的 Attention：
最终把 32 个输出沿特征维拼起来，形状回到 \([S_q,4096]\)，再乘输出投影 \(W_O\in\mathbb R^{4096\times4096}\)。这一层可通过 \(W_O\) 组合不同 head 提取的信息。
4. KV Cache 从哪里来
首次 Prefill 一次处理整个 prompt，各层把算好的 K/V 保存下来。因果约束下，过去位置在固定模型、固定前缀和位置处理下的状态无需因新 token 到来而改变。
生成时，每次把一个新 token 送入模型。到某层时，它的输入是 \([1,4096]\)。这一层只为新位置投影 Q/K/V，再把新 K/V 追加到缓存。假设追加后有 8192 个可见位置，单 head 的计算是
旧 Q 无须参加这一步。当前 query 要重新计算自己对所有可见位置的权重，因此旧 K/V 仍然要读。Prefill 最后一个位置的 logits 通常用于采样首个生成 token，下一次 Decode 再消费这个 token。
每 token、每层的 MHA 缓存元素数是 \(2H_qd_h=8192\)，乘 BF16 的 2 B，得到 16 KiB。8K 历史是一层 128 MiB；若 32 层均为相同配置，则是 4 GiB。这些值只计 KV 有效载荷。
5. MQA 将 KV head 收到一个
保持 32 个 Q head，但让所有 head 共用一个 K head、一个 V head。投影形状变成
每个 Q head 仍然有自己的 query，公式变成
这里可以直接套用前面两个 query 的手算例子。K/V 相同，分数和输出仍可不同。MQA 的每 token、每层缓存是 \(2\times1\times128\times2=512\) B，8K 历史是一层 4 MiB。相对同维度 MHA 是三十二分之一。共享减少了独立 K/V 投影的表达自由度，质量需要通过训练与评测验证。MQA 原论文
6. GQA 在多个组之间保留独立 K/V
取 8 个 KV head，每 4 个 Q head 组成一组。Q head 0–3 使用 KV head 0，4–7 使用 KV head 1，依此类推。定义
投影形状现在是
Prefill 的 Q 是 \([32,4096,128]\)，K/V 是 \([8,4096,128]\)。缓存里只需保存后者的 8 组。把 KV 复制到 32 组可以方便写参考实现，优化内核应通过索引或广播读取共享数据。
GQA 的缓存为每 token、每层 \(2\times8\times128\times2=4096\) B，即 4 KiB；8K 历史是一层 32 MiB。GQA 原论文还研究了从 MHA checkpoint 转换并继续训练的方式。已经训练好的独立 K/V 权重不能只靠运行时分组就保持输出不变。GQA 原论文
图中为便于观察取八个 Q head。正文算例使用 32 个 Q head。
7. 一组 Q 的 GEMM 形状揭示了复用机会
Decode 取出一个 KV group 的 4 行 query：
这一组可组织成
每行分别归一化。K/V 的一个 tile 只要搬进片上一次，就可以服务四行 Q。
假设 KV tile 长度取 256，QK 是 \([4,128]\times[128,256]\)，PV 是 \([4,256]\times[256,128]\)。遍历不同 tile 时需要保存每行的最大值、指数和与输出累积，最后得到整段历史的归一化输出。这部分在线 Softmax 会在后续内核一讲推导。
Prefill 里，一个 KV group 的多个 head、多个 query 位置可以共同提供更多行。例如 16 个位置乘 4 个 head 共 64 行。但每行对应的 query 位置不同，因果 mask 仍需按原位置处理。
8. 缓存节省与计算节省分别有多少
每个 query head 都计算一次 QK 与一次 PV。只看这两项，单 token Decode 合计 FLOPs 约为
代入得到 \(134217728\) FLOPs，约 134.2 MFLOPs。MHA、GQA、MQA 在这个固定配置下相同。K/V 投影的宽度会随 \(H_{kv}\) 缩小，因此投影部分的 FLOPs 确实减少。
| 结构 | 8K 历史 KV，每层 | QK+PV，单 token 每层 | 理想 KV 读取算术强度 |
|---|---|---|---|
| MHA，32 KV head | 128 MiB | 134.2 MFLOPs | 1 FLOP/B |
| GQA，8 KV head | 32 MiB | 134.2 MFLOPs | 4 FLOP/B |
| MQA，1 KV head | 4 MiB | 134.2 MFLOPs | 32 FLOP/B |
最后一列假设每份 KV 仅从 HBM 读一次，忽略权重、Q/O、中间结果和页元数据。GQA 若四个 Q head 分开运行并各自重读同一份 KV，实际读量可接近 MHA。任务划分和片上复用会决定理想收益能实现多少。
9. 放到 NPU 上继续看
对一个假设 K tile 宽为 256 的矩阵阵列，QK 的规约只有 128。如果只能补零映射这一维，则有效比例为一半。但阵列可能支持子阵列、不同精度或独立任务打包，要按具体指令形状判断。
同时，Decode 的 \(M=4\) 也很小。即使把 K 维填满，行维仍可能不足。增加活跃请求、组内联合处理 Q head、沿历史拆分工作，都能提供并行任务，不过不同请求的 KV 不同，不能直接把它们当成共享同一 K 的一个大 GEMM。
片上空间还会约束 tile。BF16 的一个 256-token KV tile，两份数据共 \(2\times256\times128\times2=128\) KiB，再加 Q、分数、输出和双缓冲。tile 取大能提升复用，也可能增加缓冲压力。实际调优先固定模型和输入，测 HBM 字节数、矩阵利用率、Vector 归约与同步时间，再决定改变哪个维度。
下一讲把共享进一步推进到 latent 表示。先让 K/V 由同一个低维向量生成，再逐步推导怎样避免每一步展开整个历史。继续阅读 MLA
