""" ============================================================================= CS349E · 第 2 讲 · 05_attention_variants.py —— MHA、GQA、MQA、MLA：为推理而改的注意力 ============================================================================= 【这个脚本要回答的问题】 04 脚本说明了：decode 受限于读 KV cache，而且批再大，注意力核心的强度也停在一个上限。 这几种注意力变体都是**为了推理**改的：在尽量不损失质量的前提下，让每个 token 的 KV 变小。 演示 1：GQA = "几个头共用一组 K/V"的 MHA：把 GQA 的 K/V 复制给每个头，得到一个输出完全相同的 MHA； KV 组数 G 从 8 降到 1，每 token 的 KV 和注意力核心的强度怎样变化（实际运行并记账）。 演示 2：把训练好的 MHA 改成 GQA（GQA 论文的"均值池化"）：头和头越像，转换的误差越小； 头各不相同时误差很大——所以论文要再"补训练"5% 的算力。 演示 3：MLA（DeepSeek-V2/V3）：只缓存一个低维的"潜在向量"，用矩阵结合律把上投影"吸收"进 query， 不用解压就能算注意力；验证两种算法输出相同，并数一数 FLOPs 和字节。为什么 RoPE 要单独留一份。 演示 4：放到 DeepSeek-V3 的注意力尺寸上：四种变体的 KV / token、读 KV 的时间、注意力的强度。 【怎么读输出】全部是确定性的计算，每次运行结果相同。 【运行】python3 05_attention_variants.py （约 2 秒） 【依赖】numpy；同目录的 transformer_lab；从 ../lecture01 导入 console、hw_specs。 """ import math import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_bytes, fmt_num, fmt_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import H100 # noqa: E402 from transformer_lab import (ATTN_CORE_OPS, TINY, Counter, KVCache, apply_rope, config_like, count_params, # noqa: E402 forward, init_params, rope_angles, softmax) # --------------------------------------------------------------------------- # 演示 1：GQA 是共用 K/V 的 MHA # --------------------------------------------------------------------------- def expand_gqa_to_mha(params, cfg): """ 把 GQA 的参数展开成 MHA：每个 query 头分到"自己那一组"K/V 投影的一份拷贝。 【思路】GQA 里第 h 个 query 头用第 h // R 组 K/V（R = H/G）。把 Wk 的第 g 组列复制 R 份，排成 H 组， 就得到一个每个头都有自己 K/V 的 MHA——只是同一组里的头恰好完全相同。 【返回】(新参数, 新配置) 【Python 小知识】np.repeat(a, R, axis=1) 把 a 沿第 1 维每个元素重复 R 次（[a, a, b, b, …]）。 """ H, G, dh = cfg.n_heads, cfg.n_kv_heads, cfg.d_head R = H // G new = {k: v for k, v in params.items() if k != "layers"} new["layers"] = [] for lp in params["layers"]: q = dict(lp) for name in ("wk", "wv"): W = lp[name].reshape(cfg.d_model, G, dh) q[name] = np.repeat(W, R, axis=1).reshape(cfg.d_model, H * dh) new["layers"].append(q) return new, config_like(cfg, n_kv_heads=H) def demo1_gqa_is_shared_mha(): print_header("演示 1：GQA = 几个头共用一组 K/V 的 MHA；KV 组数与强度") cfg = TINY p = init_params(cfg, seed=0) p_mha, cfg_mha = expand_gqa_to_mha(p, cfg) toks = np.random.default_rng(0).integers(0, cfg.vocab, size=(2, 20)) diff = np.max(np.abs(forward(p, cfg, toks) - forward(p_mha, cfg_mha, toks))) print(f"GQA（{cfg.n_kv_heads} 组 KV）与'复制 K/V 得到的 MHA'（{cfg_mha.n_kv_heads} 组）的 logits 最大差：{diff:.1e}") rows = [] B, ctx = 16, 128 for G in (8, 4, 2, 1): c_ = config_like(TINY, n_kv_heads=G) pp = init_params(c_, seed=0) rng = np.random.default_rng(1) cache = KVCache(c_, B) forward(pp, c_, rng.integers(0, c_.vocab, size=(B, ctx - 1)), cache=cache) cnt = Counter() forward(pp, c_, rng.integers(0, c_.vocab, size=(B, 1)), cache=cache, counter=cnt, last_only=True) fl = cnt.total("useful", ATTN_CORE_OPS) kv = cnt.total("kv", ATTN_CORE_OPS) name = {8: "MHA", 1: "MQA"}.get(G, f"GQA-{G}") rows.append([name, G, c_.group, f"{count_params(pp):,}", f"{c_.to_model_config().kv_bytes_per_token():.0f} B", fmt_bytes(kv), f"{fl / kv:.1f}"]) print(f"\n小模型（H = 8 个 query 头），decode：B = {B}，上下文 {ctx}，实际运行并记账：") print_table(["变体", "KV 组数 G", "组大小 H/G", "参数量", "KV / token", "这一步读 KV", "注意力核心强度"], rows) print_observations([ "把 GQA 的 K/V 投影复制给组里每个头，就得到一个输出相同（只差 1e-15 量级的舍入）的 MHA：GQA 只是'强制同组的头共用 K/V'的 MHA，" "所以它能表达的东西是 MHA 的一个子集，换来的是不必存那些重复的 K/V。", "G 每减半，KV / token 和读 KV 的字节就减半；注意力核心的强度恰好等于组大小 H/G（1 → 2 → 4 → 8 FLOP/byte）：" "一组 K/V 读进来，被组里每个 query 头用一次。", "参数量几乎不变（K、V 投影在总参数里占比很小）——GQA 省的是**推理时**的显存和带宽，而不是模型大小。", ]) # --------------------------------------------------------------------------- # 演示 2：MHA → GQA 的均值池化转换 # --------------------------------------------------------------------------- def mha_with_similar_heads(cfg_mha, groups, rho, seed=0): """ 构造一个 MHA 模型：每 H/groups 个相邻的头属于一组，组内各头的 K/V 投影相关系数为 rho。 【思路】W_head = √rho · W_组公共部分 + √(1−rho) · W_头自己的部分。rho = 1 时组内各头完全相同， rho = 0 时各头独立。真实训练出的模型里，头和头之间有一定的冗余，介于两者之间。 """ p = init_params(cfg_mha, seed=seed) rng = np.random.default_rng(seed + 100) H, dh, d = cfg_mha.n_heads, cfg_mha.d_head, cfg_mha.d_model R = H // groups for lp in p["layers"]: for name in ("wk", "wv"): shared = rng.standard_normal((d, groups, 1, dh)) / math.sqrt(d) own = rng.standard_normal((d, groups, R, dh)) / math.sqrt(d) lp[name] = (math.sqrt(rho) * shared + math.sqrt(1 - rho) * own).reshape(d, H * dh) return p def mean_pool_to_gqa(params, cfg, groups): """GQA 论文的转换：把同一组里各头的 K/V 投影矩阵取平均，得到一组 K/V。""" H, dh, d = cfg.n_heads, cfg.d_head, cfg.d_model R = H // groups new = {k: v for k, v in params.items() if k != "layers"} new["layers"] = [] for lp in params["layers"]: q = dict(lp) for name in ("wk", "wv"): q[name] = lp[name].reshape(d, groups, R, dh).mean(axis=2).reshape(d, groups * dh) new["layers"].append(q) return new, config_like(cfg, n_kv_heads=groups) def kl_and_agree(logits_a, logits_b): """两组 logits 的平均 KL(p_a ‖ p_b)（每个位置一次）和 top-1 一致率。""" pa, pb = softmax(logits_a), softmax(logits_b) kl = np.sum(pa * (np.log(pa) - np.log(pb)), axis=-1) agree = np.mean(np.argmax(logits_a, -1) == np.argmax(logits_b, -1)) return float(np.mean(kl)), float(agree) def demo2_mean_pool_conversion(): print_header("演示 2：把 MHA 改成 GQA（均值池化）——头越像，转换越准") cfg = config_like(TINY, n_kv_heads=8) toks = np.random.default_rng(2).integers(0, cfg.vocab, size=(4, 64)) rows = [] for rho in (1.0, 0.9, 0.5, 0.0): cells = [f"{rho:.1f}"] for groups in (4, 2, 1): p = mha_with_similar_heads(cfg, groups=groups, rho=rho) ref = forward(p, cfg, toks) pg, cg = mean_pool_to_gqa(p, cfg, groups) kl, ag = kl_and_agree(ref, forward(pg, cg, toks)) cells.append(f"{kl:.4f} / {ag * 100:.0f}%") rows.append(cells) print("每格：转换后与原 MHA 的平均 KL 散度 / top-1 一致率（4 个序列 × 64 个位置）") print_table(["组内头的相关系数 ρ", "→ GQA-4", "→ GQA-2", "→ MQA"], rows) print_observations([ "ρ = 1（组内各头完全相同）时转换是精确的（KL = 0、100% 一致）：这就是演示 1 的反方向。", "头越不像，均值池化丢的信息越多：每一列的 KL 都随 ρ 下降而增大；同一个 ρ 下，组越大（GQA-4 → MQA）KL 越大。" "top-1 一致率大体同向（ρ = 0 转成 MQA 时最低，66%），但有小的波动（每一列是一个'组结构和目标一致'的不同随机模型）。", "GQA 论文的做法是：先均值池化（他们比较过，比'只保留一个头'或'随机初始化'都好），再用原预训练算力的 5% 继续训练" "（uptraining）把损失补回来。论文报告 T5-XXL 转成 GQA-8 后质量接近 MHA、速度接近 MQA（推理时间 0.28 ms vs MHA 1.51 ms、" "MQA 0.24 ms）。这里没有训练，只展示'为什么需要补训练'。", ]) # --------------------------------------------------------------------------- # 演示 3：MLA # --------------------------------------------------------------------------- def make_mla(d, H, dh, d_c, d_r, seed=0): """ 一层 MLA 注意力的参数（DeepSeek-V3 技术报告 §2.1.1 的公式 1–11，省略了 query 的低秩压缩）。 【参数】d：输入维；H：头数；dh：每头维度（不含 RoPE 部分）；d_c：KV 潜在向量的维度；d_r：RoPE 部分的维度 【返回】字典： w_dkv [d, d_c] ：下投影：h → c（只缓存它） w_uk [H, d_c, dh] ：上投影：c → 每个头的 key（不带 RoPE 的部分） w_uv [H, d_c, dh] ：上投影：c → 每个头的 value w_kr [d, d_r] ：h → 共享的 RoPE key（所有头共用一份，也要缓存） w_q [d, H, dh]、w_qr [d, H, d_r]：query 的两部分 w_o [H, dh, d] ：输出投影 """ rng = np.random.default_rng(seed) def r(*shape, fan): return rng.standard_normal(shape) / math.sqrt(fan) return {"w_dkv": r(d, d_c, fan=d), "w_uk": r(H, d_c, dh, fan=d_c), "w_uv": r(H, d_c, dh, fan=d_c), "w_kr": r(d, d_r, fan=d), "w_q": r(d, H, dh, fan=d), "w_qr": r(d, H, d_r, fan=d), "w_o": r(H, dh, d, fan=H * dh)} def mla_naive(P, h): """ "解压"写法：从缓存的 c 重建每个头的完整 K、V，再做普通的多头注意力。 【参数】P：make_mla 的参数；h：[T, d]（一个序列；位置 0 … T−1；因果掩码） 【返回】[T, d] """ T = h.shape[0] H, _, dh = P["w_uk"].shape d_r = P["w_kr"].shape[1] cos, sin = rope_angles(np.arange(T), d_r) c = h @ P["w_dkv"] # [T, d_c] ← 缓存 k_rope = apply_rope((h @ P["w_kr"])[None, :, None, :], cos, sin)[0, :, 0] # [T, d_r] ← 缓存 k = np.einsum("tc,hcd->thd", c, P["w_uk"]) # 解压出 [T, H, dh] v = np.einsum("tc,hcd->thd", c, P["w_uv"]) q = np.einsum("tm,mhd->thd", h, P["w_q"]) q_rope = apply_rope(np.einsum("tm,mhr->thr", h, P["w_qr"])[None], cos, sin)[0] scores = (np.einsum("thd,shd->hts", q, k) + np.einsum("thr,sr->hts", q_rope, k_rope)) / math.sqrt(dh + d_r) scores = np.where(np.tril(np.ones((T, T), bool))[None], scores, -np.inf) o = np.einsum("hts,shd->thd", softmax(scores), v) return np.einsum("thd,hdm->tm", o, P["w_o"]) def mla_absorbed(P, h): """ "吸收"写法：不重建 K、V，直接在潜在向量 c 上做注意力。 【思路】矩阵乘法满足结合律： q_h · k_h = q_h · (c · W_uk[h])ᵀ = (q_h · W_uk[h]ᵀ) · cᵀ ← 先把 W_uk 乘进 query（"吸收"） Σ_s p_s · v_s = Σ_s p_s · (c_s · W_uv[h]) = (Σ_s p_s · c_s) · W_uv[h] ← 先在 c 上加权求和，最后再上投影 这样每一步 decode 只需读缓存里的 c（d_c 维）和共享的 RoPE key（d_r 维），不用为每个头展开 K、V。 【为什么 RoPE 部分不能吸收】RoPE 是"按位置旋转"：q 在位置 m 转 R_m，k 在位置 n 转 R_n， (R_m·q)ᵀ(R_n·W·c) = qᵀ·R_mᵀ·R_n·W·c，中间的 R_mᵀR_n 随位置变化，不能并成一个固定的矩阵。 所以 DeepSeek 把 key 拆成两部分：不带 RoPE 的部分走潜在向量（可以吸收），带 RoPE 的小部分单独算、单独缓存。 """ T = h.shape[0] H, _, dh = P["w_uk"].shape d_r = P["w_kr"].shape[1] cos, sin = rope_angles(np.arange(T), d_r) c = h @ P["w_dkv"] k_rope = apply_rope((h @ P["w_kr"])[None, :, None, :], cos, sin)[0, :, 0] q = np.einsum("tm,mhd->thd", h, P["w_q"]) q_abs = np.einsum("thd,hcd->thc", q, P["w_uk"]) # 吸收：[T, H, d_c] q_rope = apply_rope(np.einsum("tm,mhr->thr", h, P["w_qr"])[None], cos, sin)[0] scores = (np.einsum("thc,sc->hts", q_abs, c) + np.einsum("thr,sr->hts", q_rope, k_rope)) / math.sqrt(dh + d_r) scores = np.where(np.tril(np.ones((T, T), bool))[None], scores, -np.inf) o_lat = np.einsum("hts,sc->thc", softmax(scores), c) # 在潜在空间里加权求和：[T, H, d_c] o = np.einsum("thc,hcd->thd", o_lat, P["w_uv"]) # 最后再上投影 return np.einsum("thd,hdm->tm", o, P["w_o"]) def mla_decode_costs(H, dh, d_c, d_r, S, bpe=2.0): """ 一个 token、一层、上下文 S 时，MLA 三种实现的注意力 FLOPs 和读 KV 的字节（不含 Q/O 投影）。 【三种实现】 展开缓存：把解压后的 K、V（每头 dh + d_r 和 dh）存下来——字节和 MHA 一样，省不了显存 每步解压：只存 c 和 RoPE key，每一步把 S 个 c 重新解压成 K、V（多出 2·S·d_c·H·dh·2 次运算） 吸收 ：只存 c 和 RoPE key，在潜在空间里算（每个头对每个位置做 d_c + d_r 维和 d_c 维两个点积） 【返回】{名字: (FLOPs, 字节)} """ mha_like = 4 * H * dh * S + 2 * H * d_r * S return { "展开缓存（存解压后的 K/V）": (mha_like, S * H * (2 * dh + d_r) * bpe), "每步解压（存 c，每步重建 K/V）": (mha_like + 2 * 2 * S * d_c * H * dh, S * (d_c + d_r) * bpe), "吸收（存 c，在潜在空间里算）": (2 * H * (d_c + d_r) * S + 2 * H * d_c * S + 2 * 2 * H * dh * d_c, S * (d_c + d_r) * bpe), } def demo3_mla(): print_header("演示 3：MLA——只缓存潜在向量，用结合律'吸收'上投影") d, H, dh, d_c, d_r = 128, 8, 16, 32, 8 P = make_mla(d, H, dh, d_c, d_r) h = np.random.default_rng(3).standard_normal((24, d)) diff = np.max(np.abs(mla_naive(P, h) - mla_absorbed(P, h))) print(f"小例子（d = {d}，H = {H}，dh = {dh}，d_c = {d_c}，d_r = {d_r}，24 个 token）：解压写法 vs 吸收写法 最大差 {diff:.1e}") mha = 2 * H * dh print(f"每 token 每层要缓存的数：MHA {mha}，MLA {d_c} + {d_r} = {d_c + d_r}（{mha / (d_c + d_r):.1f} 倍小）") H, dh, d_c, d_r, S = 128, 128, 512, 64, 32768 rows = [] for name, (fl, by) in mla_decode_costs(H, dh, d_c, d_r, S).items(): rows.append([name, fmt_num(fl), fmt_bytes(by), f"{fl / by:.0f}"]) print(f"\nDeepSeek-V3 的注意力尺寸（H = {H}，dh = {dh}，d_c = {d_c}，d_r = {d_r}），一个 token、一层、上下文 {S}：") print_table(["实现", "注意力 FLOPs", "读 KV 字节", "强度 FLOP/byte"], rows) print_observations([ "两种写法的输出相同（只差浮点舍入）：吸收只是利用矩阵乘法的结合律换了一个计算顺序。", "'展开缓存'省不了任何显存（和 MHA 一样大）；'每步解压'省了显存，却要每一步把整个上下文重新解压，运算量是展开缓存的约 400 倍；" "'吸收'两头都占：只读 576 个数 / token / 层，运算量约是展开缓存的 3.4 倍。", "吸收后的强度约 240 FLOP/byte，接近 H100 的屋脊点 295：MLA 把 decode 注意力从'严重 memory-bound'推到了接近平衡。" "这是用 FLOPs 换字节——在算力富余、带宽紧张的 decode 阶段，这笔交易很划算。", "代价：带 RoPE 的那 64 维不能吸收，要单独缓存；而且 d_c = 512 是所有头共享的，第 4 讲按头切分张量并行时，" "MLA 的潜在向量切不开（每块卡都要存一整份），这让 MLA 更适合按数据 / 批切分的注意力。", ]) # --------------------------------------------------------------------------- # 演示 4：DeepSeek-V3 尺寸上的四种变体 # --------------------------------------------------------------------------- def demo4_at_scale(): print_header("演示 4：同一个注意力尺寸（DeepSeek-V3：61 层、128 头、dh = 128）下的四种变体") L, H, dh, d_c, d_r = 61, 128, 128, 512, 64 ctx, B = 32768, 32 chip = H100 variants = [] for name, G in (("MHA", 128), ("GQA-8", 8), ("MQA", 1)): kv_tok = 2 * L * G * dh * 2 fl_tok = L * 4 * H * dh * ctx variants.append((name, kv_tok, fl_tok)) kv_mla = L * (d_c + d_r) * 2 fl_mla = L * (2 * H * (d_c + d_r) * ctx + 2 * H * d_c * ctx + 4 * H * dh * d_c) variants.append(("MLA（吸收）", kv_mla, fl_mla)) rows = [] for name, kv_tok, fl_tok in variants: by = B * ctx * kv_tok fl = B * fl_tok t = max(fl / chip.peak_flops("bf16"), by / chip.mem_bw) rows.append([name, fmt_bytes(kv_tok, binary=True), fmt_bytes(ctx * kv_tok), f"{fl / by:.0f}", fmt_time(by / chip.mem_bw), fmt_time(fl / chip.peak_flops("bf16")), fmt_time(t)]) print(f"BF16 KV，上下文 {ctx}，批大小 {B}，一块 H100 上一步 decode 的注意力部分（不含投影和 MLP）：") print_table(["变体", "KV / token", "一个 32K 序列的 KV", "强度", "读 KV 时间", "计算时间", "下界"], rows) print_observations([ f"MHA 每个 token 要 {fmt_bytes(variants[0][1], binary=True)}：一个 32K 的序列就要 {fmt_bytes(ctx * variants[0][1])}，" "一块 H100 放不下一个序列的 KV。DeepSeek-V3 这么宽的注意力，不压缩 KV 根本没法 decode。", "GQA-8 和 MQA 把 KV 缩小 16 倍和 128 倍；MLA 每 token 只要 576 个数 / 层，大小介于 GQA-8 和 MQA 之间" "（相当于 576 / 256 = 2.25 组 KV 头的 GQA），DeepSeek 报告它的质量和 MHA 相当。", "MLA 的读 KV 时间和计算时间几乎相等（强度约 240，接近屋脊点）：它是四种里唯一让 decode 注意力'算和搬都不闲着'的。" "代价是计算时间是其他三种的 4 倍多（18 ms vs 4.2 ms）——在更快的低精度算力上（FP8），这部分会再缩短。", "【系统联系】这是一个**模型结构为硬件让步**的例子：GQA、MQA、MLA 都不是为了更好的质量，而是为了 decode 的带宽和显存。" "下一讲的量化（把 KV 存成 FP8 / FP4）可以在此基础上再省一半到四分之三。", ]) if __name__ == "__main__": demo1_gqa_is_shared_mha() demo2_mean_pool_conversion() demo3_mla() demo4_at_scale()