Causal Self-Attention

Attention 精讲

用图把 mask、多头、GQA、KV Cache 钉死。公式用等宽文本，避免 Markdown 下划线坑。

完整推导亦见 notes/01b_Attention精读.md · 跑 code/01b_*.py

1. 一条链

x → Wq/Wk/Wv → Q,K,V
  → (RoPE)
  → S = (Q @ Kᵀ) / √d
  → causal mask（未来 = -∞）
  → A = softmax(S)      # 每行和为 1
  → Y = A @ V
  → concat heads → Wo

一句话：用 query–key 相似度当权重，对 value 做加权平均；因果约束让位置 t 看不到未来。

2. 因果 Mask 长什么样

因果 mask 矩阵 允许看（保留分数） 屏蔽（置 -∞） t=0 t=1 t=2 t=3 j=0 1 2 3 下三角（含对角）= 可见 上三角 = 未来，softmax 前抹掉 PyTorch: triu(..., diagonal=1) masked_fill(-inf)
图 1 · 因果 mask：行是 query 位置 t，列是 key 位置 j
Mask 必须加在 softmax 之前。加在之后会破坏概率归一，且 -inf 技巧就是为 softmax 准备的。

3. 多头与 GQA

MHA vs GQA MHA nh 个 Q 头 nh 个 K 头 nh 个 V 头 KV Cache ∝ nh GQA nh 个 Q 头 nkv 个 K/V 头 计算时 repeat 对齐 KV Cache ∝ nkv 体积比 = nkv/nh
图 2 · GQA 用更少 KV 头省 cache；Llama2-70B 常见 nh=64, nkv=8 → 约 1/8

4. 对照代码：wq / wk / wv / wo 是什么

跑 python code/01_llama2_decoder.py 会打印一层里的 Linear。默认配置 C=256, nh=8, nkv=4, d=32 时大致是：

attn.wq  weight (256, 256)   # → Q，8 头 × 32
attn.wk  weight (128, 256)   # → K，4 头 × 32（GQA）
attn.wv  weight (128, 256)   # → V，同上
attn.wo  weight (256, 256)   # 多头拼回后再投影

nn.Linear 的 weight 形状是 (out_features, in_features)。

名字人话算出什么
wq提问怎么组织Q＝「我要找什么」
wk标签怎么挂K＝「我有什么可被查找」
wv货架上的货V＝真正被加权混合的内容
wo汇总翻译多头结果拼回 C 维，写回残差流
 
Q = x @ Wqᵀ ;  K = x @ Wkᵀ ;  V = x @ Wvᵀ
分数 ≈ (Q @ Kᵀ) / √d   → mask → softmax → 加权 V
out = Y @ Woᵀ
为何 wk/wv 是 128 不是 256？ GQA：Q 有 8 头，K/V 只存 4 头；计算时把每个 KV 头 repeat 给 2 个 Q 头。 把 n_kv_heads 改成 8 再跑，wk/wv 就会变成 (256, 256)（普通 MHA）。

同一脚本里的 MLP（SwiGLU）也会打印：

mlp.w1  (688, 256)   # gate
mlp.w3  (688, 256)   # up
mlp.w2  (256, 688)   # down
# out = W2( SiLU(W1 x) * (W3 x) )

结构页有 Decoder 整层对照；本页抓住：四个 W 就是把 x 变成 Q/K/V 再投影回去的可学习矩阵。

5. KV Cache 字节（一层、fp16）

bytes_MHA = 2 * B * T * nh * d * 2
          = 4 * B * T * C          # 因 nh*d = C

bytes_GQA = 4 * B * T * C * (nkv / nh)

例：B=1, T=1024, C=4096 → MHA 一层约 16 MiB；GQA nkv=nh/8 → 2 MiB。

GPT-3 量级 B=16, T=4096, C=12288 → 一层 MHA 约 3 GiB。8MB SRAM 塞不下 1 层——这就是飞书题要的数量级直觉。

6. Prefill vs Decode

阶段Attention 算力直觉常像
Prefill 长度 T∝ T² · C算力 / IO 双吃
Decode 一步（cache 长 L）∝ L · C读 KV → 带宽瓶颈

这和 NPU/编译器主线直接相关：Prefill 要 tiling / 融合；Decode 要减 KV、布局、量化。

7. 动手

  1. python code/01b_attention_numeric_walkthrough.py — 看 T=3 mask
  2. 跑 01_llama2_decoder.py，对照本节 wq/wk/wv/wo；再改 n_kv_heads
  3. 默画图 1 的下三角，口述「为何 mask 在 softmax 前」

← GPT 结构 · 下一篇：MLA →